Пример #1
0
def energy_minimization(item, platform_name='CUDA', verbose=True):

    from simtk.openmm import LangevinIntegrator, Platform, Context, LocalEnergyMinimizer_minimize
    from simtk import unit

    # Integrator.

    integrator = LangevinIntegrator(0 * unit.kelvin, 1.0 / unit.picoseconds,
                                    2.0 * unit.femtoseconds)

    # Platform.

    platform = Platform.getPlatformByName(platform_name)

    # Context.

    context = Context(item.system, integrator, platform)
    context.setPositions(item.coordinates)

    # Minimization.

    if verbose == True:
        energy = context.getState(getEnergy=True).getPotentialEnergy()
        print('Potential energy before minimization: {}'.format(energy))

    LocalEnergyMinimizer_minimize(context)

    if verbose == True:
        energy = context.getState(getEnergy=True).getPotentialEnergy()
        print('Potential energy after minimization: {}'.format(energy))

    item.coordinates = context.getState(getPositions=True).getPositions(
        asNumpy=True)

    pass
Пример #2
0
def get_potential_energy(item, coordinates=None, platform_name='CUDA'):

    from simtk.openmm import LangevinIntegrator, Platform, Context
    from simtk import unit
    import numpy as np

    integrator = LangevinIntegrator(0.0 * unit.kelvin, 0.0 / unit.picoseconds,
                                    2.0 * unit.femtoseconds)

    platform = Platform.getPlatformByName(platform_name)

    context = Context(item.system, integrator, platform)

    if coordinates is None:
        context.setPositions(item.coordinates)
    else:
        context.setPositions(coordinates)

    if item.box is not None:
        context.setPeriodicBoxVectors(item.box[0], item.box[1], item.box[2])

    state = context.getState(getEnergy=True)
    potential_energy = state.getPotentialEnergy()

    return potential_energy
Пример #3
0
class OpenMMStateBuilder(StateBuilder):
    """Build an OpenMM "state" that can be sent to a device to simulate.
    """
    def __init__(self, system, integrator=None):

        # if strings are passed in, assume that they are paths to
        # xml files on disk
        if isinstance(system, basestring):
            with open(system) as f:
                system = XmlSerializer.deserialize(f.read())
        if isinstance(integrator, basestring):
            with open(integrator) as f:
                integrator = XmlSerializer.deserialize(f.read())

        if integrator is None:
            # this integrator isn't really necessary, but it has to be something
            # for the openmm API to let us serialize the state
            integrator = VerletIntegrator(2 * femtoseconds)
        self.context = Context(system, integrator,
                               Platform.getPlatformByName('Reference'))

    def build(self, trajectory):
        """Create a serialized XML state from the first frame in a trajectory

        Parameteters
        ------------
        trajectory : mdtraj.trajectory.Trajectory
            The trajectory to take the frame from. We'll use both the the
            positions and the box vectors (if you're using periodic boundary
            conditions)
        """
        periodic = False
        if trajectory.unitcell_vectors is not None:
            a, b, c = trajectory.unitcell_lengths[0]
            np.testing.assert_array_almost_equal(trajectory.unitcell_angles[0],
                                                 np.ones(3) * 90)
            self.context.setPeriodicBoxVectors([a, 0, 0] * nanometers,
                                               [0, b, 0] * nanometers,
                                               [0, 0, c] * nanometers)
            periodic = True

        self.context.setPositions(trajectory.openmm_positions(0))
        state = self.context.getState(getPositions=True,
                                      getVelocities=True,
                                      getForces=True,
                                      getEnergy=True,
                                      getParameters=True,
                                      enforcePeriodicBox=periodic)
        return XmlSerializer.serialize(state)
Пример #4
0
class OpenMMStateBuilder(StateBuilder):
    """Build an OpenMM "state" that can be sent to a device to simulate.
    """
    def __init__(self, system, integrator=None):

        # if strings are passed in, assume that they are paths to
        # xml files on disk
        if isinstance(system, basestring):
            with open(system) as f:
                system = XmlSerializer.deserialize(f.read())
        if isinstance(integrator, basestring):
            with open(integrator) as f:
                integrator = XmlSerializer.deserialize(f.read())

        if integrator is None:
            # this integrator isn't really necessary, but it has to be something
            # for the openmm API to let us serialize the state
            integrator = VerletIntegrator(2*femtoseconds)
        self.context = Context(system, integrator, Platform.getPlatformByName('Reference'))

    def build(self, trajectory):
        """Create a serialized XML state from the first frame in a trajectory

        Parameteters
        ------------
        trajectory : mdtraj.trajectory.Trajectory
            The trajectory to take the frame from. We'll use both the the
            positions and the box vectors (if you're using periodic boundary
            conditions)
        """
        periodic = False
        if trajectory.unitcell_vectors is not None:
            a, b, c = trajectory.unitcell_lengths[0]
            np.testing.assert_array_almost_equal(trajectory.unitcell_angles[0], np.ones(3)*90)
            self.context.setPeriodicBoxVectors([a, 0, 0] * nanometers, [0, b, 0] * nanometers, [0, 0, c] * nanometers)
            periodic = True

        self.context.setPositions(trajectory.openmm_positions(0))
        state = self.context.getState(getPositions=True, getVelocities=True,
                                      getForces=True, getEnergy=True,
                                      getParameters=True, enforcePeriodicBox=periodic)
        return XmlSerializer.serialize(state)
Пример #5
0
class Langevin():

    _explorer = None
    _initialized = False
    _context = None
    _integrator = None

    _timestep = Quantity(value=2.0, unit=u.femtosecond)
    _temperature = Quantity(value=298.0, unit=u.kelvin)
    _collision_rate = Quantity(value=1.0, unit=1.0 / u.picosecond)

    def __init__(self, explorer):

        self._explorer = explorer

    def _initialize(self):

        system = self._explorer.context.getSystem()
        platform = self._explorer.context.getPlatform()
        properties = {}
        if platform.getName() == 'CUDA':
            properties['CudaPrecision'] = 'mixed'

        self._integrator = LangevinIntegrator(self._temperature,
                                              self._collision_rate,
                                              self._timestep)
        self._context = Context(system, self._integrator, platform, properties)
        self._initialized = True

    def set_parameters(self,
                       temperature=Quantity(value=298.0, unit=u.kelvin),
                       collision_rate=Quantity(value=1.0,
                                               unit=1.0 / u.picosecond),
                       timestep=Quantity(value=2.0, unit=u.femtosecond)):

        self._timestep = timestep
        self._temperature = temperature
        self._collision_rate = collision_rate

        if self._initialized:

            self._integrator.setFriction(
                self._collision_rate.value_in_unit(u.picosecond**-1))
            self._integrator.setTemperature(
                self._temperature.value_in_unit(u.kelvin))
            self._integrator.setStepSize(
                self._timestep.value_in_unit(u.picoseconds))

        else:

            self._initialize()

    def get_parameters(self):

        parameters = {
            'timestep': self._timestep,
            'temperature': self._temperature,
            'collision_rate': self._collision_rate
        }

        return parameters

    def replicate_parameters(self, explorer):

        timestep = explorer.md.langevin._timestep
        temperature = explorer.md.langevin._temperature
        collision_rate = explorer.md.langevin._collision_rate

        self.set_paramters(temperature, collision_rate, timestep)

    def _set_coordinates(self, coordinates):

        self._context.setPositions(coordinates)

    def _get_coordinates(self):

        return self._context.getState(getPositions=True).getPositions(
            asNumpy=True)

    def _set_velocities(self, velocities):

        self._context.setVelocities(velocities)

    def _get_velocities(self):

        return self._context.getState(getVelocities=True).getVelocities(
            asNumpy=True)

    def _coordinates_to_explorer(self):

        self._explorer.set_coordinates(self._get_coordinates())

    def _coordinates_from_explorer(self):

        self._set_coordinates(self._explorer.get_coordinates())

    def _velocities_to_explorer(self):

        self._explorer.set_velocities(self._get_velocities())

    def _velocities_from_explorer(self):

        self._set_velocities(self._explorer.get_velocities())

    def get_time(self):

        return self._context.getState().getTime()

    def run(self, steps=0):

        if not self._initialized:

            self._initialize()

        self._coordinates_from_explorer()
        self._velocities_from_explorer()
        self._integrator.step(steps)
        self._coordinates_to_explorer()
        self._velocities_to_explorer()

    def __call__(self, *args, **kwargs):

        return self.run(*args, **kwargs)
Пример #6
0
class FIRE():

    _explorer = None
    _initialized = False
    _context = None
    _integrator = None

    _timestep = Quantity(0.001, u.picoseconds)
    _tolerance = None
    _alpha = 0.1
    _dt_max = Quantity(0.01, u.picoseconds)
    _f_inc = 1.1
    _f_dec = 0.5
    _f_alpha = 0.99
    _N_min = 5

    def __init__(self, explorer):

        self._explorer = explorer

    def _initialize(self):

        system = self._explorer.context.getSystem()
        platform = self._explorer.context.getPlatform()
        properties = {}
        if platform.getName() == 'CUDA':
            properties['CudaPrecision'] = 'mixed'

        self._integrator = FIREMinimizationIntegrator(
            timestep=self._timestep,
            tolerance=self._tolerance,
            alpha=self._alpha,
            dt_max=self._dt_max,
            f_inc=self._f_inc,
            f_dec=self._f_dec,
            f_alpha=self._f_alpha,
            N_min=self._N_min)

        self._context = Context(system, self._integrator, platform, properties)

        self._initialized = True

    def set_parameters(self,
                       timestep=Quantity(1.0, u.femtoseconds),
                       tolerance=None,
                       alpha=0.1,
                       dt_max=Quantity(10.0, u.femtoseconds),
                       f_inc=1.1,
                       f_dec=0.5,
                       f_alpha=0.99,
                       N_min=5):

        self._timestep = timestep.in_units_of(u.picoseconds)
        self._tolerance = tolerance
        self._alpha = alpha
        self._dt_max = dt_max.in_units_of(u.picoseconds)
        self._f_inc = f_inc
        self._f_dec = f_dec
        self._f_alpha = f_alpha
        self._N_min = N_min

        self._initialize()

    def replicate_parameters(self, explorer):

        timestep = explorer.quench.fire._timestep
        tolerance = explorer.quench.fire._tolerance
        alpha = explorer.quench.fire._alpha
        dt_max = explorer.quench.fire._dt_max
        f_inc = explorer.quench.fire._f_inc
        f_dec = explorer.quench.fire._f_dec
        f_alpha = explorer.quench.fire._f_alpha
        N_min = explorer.quench.fire._N_min

        self.set_parameters(timestep, tolerance, alpha, dt_max, f_inc, f_dec,
                            f_alpha, N_min)

    def _set_coordinates(self, coordinates):

        self._context.setPositions(coordinates)

    def _get_coordinates(self):

        return self._context.getState(getPositions=True).getPositions(
            asNumpy=True)

    def _coordinates_to_explorer(self):

        self._explorer.set_coordinates(self._get_coordinates())

    def _coordinates_from_explorer(self):

        self._set_coordinates(self._explorer.get_coordinates())

    def run(self, steps=0):

        if not self._initialized:

            self._initialize()

        self._coordinates_from_explorer()
        try:
            if steps == 0:
                while self._integrator.getGlobalVariableByName(
                        'converged') < 1:
                    self._integrator.step(50)
            else:
                self._integrator.step(steps)
            self._coordinates_to_explorer()
        except Exception as e:
            if str(e) == 'Particle coordinate is nan':
                print(
                    'NaN encountered in FIRE minimizer; falling back to L-BFGS after resetting positions'
                )
                self._explorer.quench.l_bfgs()
            else:
                raise e

        self._initialize()

    def __call__(self, *args, **kwargs):

        return self.run(*args, **kwargs)
Пример #7
0
class GradientDescent():

    _explorer = None
    _initialized = False
    _context = None
    _integrator = None

    _initial_step_size = Quantity(0.01, u.angstroms)
    _tolerance = Quantity(1.0, u.kilojoules_per_mole)

    def __init__(self, explorer):

        self._explorer=explorer

    def _initialize(self):

        system = self._explorer.context.getSystem()
        platform = self._explorer.context.getPlatform()
        properties = {}
        if platform.getName()=='CUDA':
            properties['CudaPrecision'] = 'mixed'

        self._integrator = GradientDescentMinimizationIntegrator(initial_step_size=self._initial_step_size)
        self._context = Context(system, self._integrator, platform, properties)
        self._initialized = True

    def set_parameters(self, tolerance=Quantity(1.0, u.kilojoules_per_mole),
            initial_step_size=Quantity(0.01, u.angstroms)):

        if not self._initialized:
            self._initialize()

        self._tolerance = tolerance.in_units_of(u.kilojoules_per_mole)
        self._initial_step_size = initial_step_size.in_units_of(u.nanometers)
        self._integrator.setGlobalVariableByName('step_size', self._initial_step_size._value)

    def replicate_parameters(self, explorer):

        tolerance = explorer.quench.gradient_descent._tolerance
        initial_step_size = explorer.quench.gradient_descent._initial_step_size

        self.set_parameters(tolerance, initial_step_size)

    def _set_coordinates(self, coordinates):

        self._context.setPositions(coordinates)

    def _get_coordinates(self):

        return self._context.getState(getPositions=True).getPositions(asNumpy=True)

    def _coordinates_to_explorer(self):

        self._explorer.set_coordinates(self._get_coordinates())

    def _coordinates_from_explorer(self):

        self._set_coordinates(self._explorer.get_coordinates())

    def run(self, steps=0):

        if not self._initialized:

            self._initialize()

        self._coordinates_from_explorer()

        try:
            if steps == 0:
                delta=np.infty
                while delta > self._tolerance._value:
                    self._integrator.step(50)
                    delta=self._integrator.getGlobalVariableByName('delta_energy')
            else:
                integrator.step(steps)

            self._coordinates_to_explorer()
        except Exception as e:
            if str(e) == 'Particle coordinate is nan':
                print('NaN encountered in gradient descent minimizer; falling back to L-BFGS after resetting positions')
                self._explorer.quench.l_bfgs()
            else:
                raise e

        self._initialize()

    def __call__(self, *args, **kwargs):

        return self.run(*args, **kwargs)
Пример #8
0
class Explorer():

    topology = None
    context = None
    pbc = False
    n_atoms = 0
    n_dof = 0
    md = None
    quench = None
    move = None

    def __init__(self, topology=None, system=None, pbc=False, platform='CUDA'):

        from .md import MD
        from .quench import Quench
        from .move import Move
        from .distance import Distance
        from .acceptance import Acceptance

        if topology is None:
            raise ValueError('topology is needed')

        if system is None:
            raise ValueError('system is needed')

        integrator = LangevinIntegrator(0 * u.kelvin, 1.0 / u.picoseconds,
                                        2.0 * u.femtoseconds)
        #integrator.setConstraintTolerance(0.00001)

        if platform == 'CUDA':
            platform = Platform.getPlatformByName('CUDA')
            properties = {'CudaPrecision': 'mixed'}
        elif platform == 'CPU':
            platform = Platform.getPlatformByName('CPU')
            properties = {}

        self.topology = topology
        self.context = Context(system, integrator, platform, properties)
        self.n_atoms = msm.get(self.context, target='system', n_atoms=True)

        self.n_dof = 0
        for i in range(system.getNumParticles()):
            if system.getParticleMass(i) > 0 * u.dalton:
                self.n_dof += 3
        for i in range(system.getNumConstraints()):
            p1, p2, distance = system.getConstraintParameters(i)
            if system.getParticleMass(
                    p1) > 0 * u.dalton or system.getParticleMass(
                        p2) > 0 * u.dalton:
                self.n_dof -= 1
        if any(
                type(system.getForce(i)) == CMMotionRemover
                for i in range(system.getNumForces())):
            self.n_dof -= 3

        self.pbc = pbc

        if self.pbc:
            raise NotImplementedError

        self.md = MD(self)
        self.quench = Quench(self)
        self.move = Move(self)
        self.distance = Distance(self)
        self.acceptance = Acceptance(self)

    def _copy(self):

        topology = self.topology
        coordinates = self.get_coordinates()
        system = self.context.getSystem()
        platform = self.context.getPlatform().getName()
        pbc = self.pbc

        tmp_explorer = Explorer(topology, system, pbc, platform)
        tmp_explorer.set_coordinates(coordinates)

        for ii, jj in vars(tmp_explorer.md).items():
            if not ii.startswith('_'):
                if jj._initialized:
                    jj.replicate_parameters(self)

        for ii, jj in vars(tmp_explorer.quench).items():
            if not ii.startswith('_'):
                if jj._initialized:
                    jj.replicate_parameters(self)

        for ii, jj in vars(tmp_explorer.move).items():
            if not ii.startswith('_'):
                if jj._initialized:
                    jj.replicate_parameters(self)

        return tmp_explorer

    def replicate(self, times=1):

        from copy import deepcopy

        if times == 1:

            output = self._copy()

        else:

            output = [self._copy() for ii in range(times)]

        return output

    def set_coordinates(self, coordinates):

        self.context.setPositions(coordinates)

    def get_coordinates(self):

        return self.context.getState(getPositions=True).getPositions(
            asNumpy=True)

    def set_velocities(self, velocities):

        self.context.setVelocities(velocities)

    def set_velocities_to_temperature(self, temperature):

        self.context.setVelocitiesToTemperature(temperature)

    def get_velocities(self):

        return self.context.getState(getVelocities=True).getVelocities(
            asNumpy=True)

    def get_temperature(self):

        return (2 * self.context.getState(getEnergy=True).getKineticEnergy() /
                (self.n_dof * u.MOLAR_GAS_CONSTANT_R)).in_units_of(u.kelvin)

    def get_potential_energy(self):

        energy = self.context.getState(getEnergy=True).getPotentialEnergy()
        return energy

    def get_potential_energy_gradient(self):

        gradient = -self.context.getState(getForces=True).getForces(
            asNumpy=True)
        gradient = gradient.ravel() * gradient.unit
        return gradient

    def get_potential_energy_hessian(self,
                                     mass_weighted=False,
                                     symmetric=True):
        """OpenMM single frame hessian evaluation
        Since OpenMM doesnot provide a Hessian evaluation method, we used finite difference on forces

        from: https://leeping.github.io/forcebalance/doc/html/api/openmmio_8py_source.html

        Returns
        -------
        hessian: np.array with shape 3N x 3N, N = number of "real" atoms
            The result hessian matrix.
            The row indices are fx0, fy0, fz0, fx1, fy1, ...
            The column indices are x0, y0, z0, x1, y1, ..
            The unit is kilojoule / (nanometer^2 * mole * dalton) => 10^24 s^-2
        """

        n_dof = self.n_atoms * 3
        pos = self.get_coordinates()
        hessian = np.empty(
            (n_dof, n_dof),
            dtype=float) * u.kilojoules_per_mole / (u.nanometers**2)
        # finite difference step size
        diff = 0.0001 * u.nanometer
        coef = 1.0 / (2.0 * diff)  # 1/2h

        for i in range(self.n_atoms):
            # loop over the x, y, z coordinates
            for j in range(3):
                # plus perturbation
                pos[i][j] += diff
                self.set_coordinates(pos)
                grad_plus = self.get_potential_energy_gradient()
                # minus perturbation
                pos[i][j] -= 2 * diff
                self.set_coordinates(pos)
                grad_minus = self.get_potential_energy_gradient()
                # set the perturbation back to zero
                pos[i][j] += diff
                # fill one row of the hessian matrix
                hessian[i * 3 + j] = (grad_plus - grad_minus) * coef

        if mass_weighted:
            mass = np.array([
                self.context.getSystem().getParticleMass(k).value_in_unit(
                    u.dalton) for k in range(self.n_atoms)
            ]) * u.dalton
            mass_weight = 1.0 / np.sqrt(mass) * (mass.unit**-0.5)
            mass_weight = np.repeat(mass_weight, 3) * mass_weight.unit
            hessian = np.multiply(
                hessian, mass_weight) * hessian.unit * mass_weight.unit
            hessian = np.multiply(
                hessian,
                mass_weight[:, np.newaxis]) * hessian.unit * mass_weight.unit

        # make hessian symmetric by averaging upper right and lower left
        if symmetric:
            hessian += hessian.T * hessian.unit
            hessian *= 0.5

        # recover the original position
        self.set_coordinates(pos)
        return hessian
Пример #9
0
    def addHydrogens(self, forcefield, pH=7.0, variants=None, platform=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
        variants differ in the presence or absence of particular hydrogens.  In particular, the following variants
        are supported:

        Aspartic acid:
            ASH: Neutral form with a hydrogen on one of the delta oxygens
            ASP: Negatively charged form without a hydrogen on either delta oxygen

        Cysteine:
            CYS: Neutral form with a hydrogen on the sulfur
            CYX: No hydrogen on the sulfur (either negatively charged, or part of a disulfide bond)

        Glutamic acid:
            GLH: Neutral form with a hydrogen on one of the epsilon oxygens
            GLU: Negatively charged form without a hydrogen on either epsilon oxygen

        Histidine:
            HID: Neutral form with a hydrogen on the ND1 atom
            HIE: Neutral form with a hydrogen on the NE2 atom
            HIP: Positively charged form with hydrogens on both ND1 and NE2

        Lysine:
            LYN: Neutral form with two hydrogens on the zeta nitrogen
            LYS: Positively charged form with three hydrogens on the zeta nitrogen

        The variant to use for each residue is determined by the following rules:

        1. The most common variant at the specified pH is selected.
        2. Any Cysteine that participates in a disulfide bond uses the CYX variant regardless of pH.
        3. For a neutral Histidine residue, the HID or HIE variant is selected based on which one forms a better hydrogen bond.

        You can override these rules by explicitly specifying a variant for any residue.  Also keep in mind that this
        function will only add hydrogens.  It will never remove ones that are already present in the model, regardless
        of the specified pH.

        Definitions for standard amino acids and nucleotides are built in.  You can call loadHydrogenDefinitions() to load
        additional definitions for other residue types.

        Parameters:
         - forcefield (ForceField) the ForceField to use for determining the positions of hydrogens
         - pH (float=7.0) the pH based on which to select variants
         - variants (list=None) an optional list of variants to use.  If this is specified, its length must equal the number
           of residues in the model.  variants[i] is the name of the variant to use for residue i (indexed starting at 0).
           If an element is None, the standard rules will be followed to select a variant for that residue.
         - platform (Platform=None) the Platform to use when computing the hydrogen atom positions.  If this is None,
           the default Platform will be used.
        Returns: a list of what variant was actually selected for each residue, in the same format as the variants parameter
        """
        # Check the list of variants.

        residues = list(self.topology.residues())
        if variants is not None:
            if len(variants) != len(residues):
                raise ValueError("The length of the variants list must equal the number of residues")
        else:
            variants = [None]*len(residues)
        actualVariants = [None]*len(residues)

        # Load the residue specifications.

        if not Modeller._hasLoadedStandardHydrogens:
            Modeller.loadHydrogenDefinitions(os.path.join(os.path.dirname(__file__), 'data', 'hydrogens.xml'))

        # Make a list of atoms bonded to each atom.

        bonded = {}
        for atom in self.topology.atoms():
            bonded[atom] = []
        for atom1, atom2 in self.topology.bonds():
            bonded[atom1].append(atom2)
            bonded[atom2].append(atom1)

        # Define a function that decides whether a set of atoms form a hydrogen bond, using fairly tolerant criteria.

        def isHbond(d, h, a):
            if norm(d-a) > 0.35*nanometer:
                return False
            deltaDH = h-d
            deltaHA = a-h
            deltaDH /= norm(deltaDH)
            deltaHA /= norm(deltaHA)
            return acos(dot(deltaDH, deltaHA)) < 50*degree

        # Loop over residues.

        newTopology = Topology()
        newTopology.setUnitCellDimensions(deepcopy(self.topology.getUnitCellDimensions()))
        newAtoms = {}
        newPositions = []*nanometer
        newIndices = []
        acceptors = [atom for atom in self.topology.atoms() if atom.element in (elem.oxygen, elem.nitrogen)]
        for chain in self.topology.chains():
            newChain = newTopology.addChain()
            for residue in chain.residues():
                newResidue = newTopology.addResidue(residue.name, newChain)
                isNTerminal = (residue == chain._residues[0])
                isCTerminal = (residue == chain._residues[-1])
                if residue.name in Modeller._residueHydrogens:
                    # Add hydrogens.  First select which variant to use.

                    spec = Modeller._residueHydrogens[residue.name]
                    variant = variants[residue.index]
                    if variant is None:
                        if residue.name == 'CYS':
                            # If this is part of a disulfide, use CYX.

                            sulfur = [atom for atom in residue.atoms() if atom.element == elem.sulfur]
                            if len(sulfur) == 1 and any((atom.residue != residue for atom in bonded[sulfur[0]])):
                                variant = 'CYX'
                        if residue.name == 'HIS' and pH > 6.5:
                            # See if either nitrogen already has a hydrogen attached.

                            nd1 = [atom for atom in residue.atoms() if atom.name == 'ND1']
                            ne2 = [atom for atom in residue.atoms() if atom.name == 'NE2']
                            if len(nd1) != 1 or len(ne2) != 1:
                                raise ValueError('HIS residue (%d) has the wrong set of atoms' % residue.index)
                            nd1 = nd1[0]
                            ne2 = ne2[0]
                            nd1HasHydrogen = any((atom.element == elem.hydrogen for atom in bonded[nd1]))
                            ne2HasHydrogen = any((atom.element == elem.hydrogen for atom in bonded[ne2]))
                            if nd1HasHydrogen and ne2HasHydrogen:
                                variant = 'HIP'
                            elif nd1HasHydrogen:
                                variant = 'HID'
                            elif ne2HasHydrogen:
                                variant = 'HIE'
                            else:
                                # Estimate the hydrogen positions.

                                nd1Pos = self.positions[nd1.index]
                                ne2Pos = self.positions[ne2.index]
                                hd1Delta = Vec3(0, 0, 0)*nanometer
                                for other in bonded[nd1]:
                                    hd1Delta += nd1Pos-self.positions[other.index]
                                hd1Delta *= 0.1*nanometer/norm(hd1Delta)
                                hd1Pos = nd1Pos+hd1Delta
                                he2Delta = Vec3(0, 0, 0)*nanometer
                                for other in bonded[ne2]:
                                    he2Delta += ne2Pos-self.positions[other.index]
                                he2Delta *= 0.1*nanometer/norm(he2Delta)
                                he2Pos = ne2Pos+he2Delta

                                # See whether either hydrogen would form a hydrogen bond.

                                nd1IsBonded = False
                                ne2IsBonded = False
                                for acceptor in acceptors:
                                    if acceptor.residue != residue:
                                        acceptorPos = self.positions[acceptor.index]
                                        if isHbond(nd1Pos, hd1Pos, acceptorPos):
                                            nd1IsBonded = True
                                            break
                                        if isHbond(ne2Pos, he2Pos, acceptorPos):
                                            ne2IsBonded = True
                                if ne2IsBonded and not nd1IsBonded:
                                    variant = 'HIE'
                                else:
                                    variant = 'HID'
                        elif residue.name == 'HIS':
                            variant = 'HIP'
                    if variant is not None and variant not in spec.variants:
                        raise ValueError('Illegal variant for %s residue: %s' % (residue.name, variant))
                    actualVariants[residue.index] = variant

                    # Make a list of hydrogens that should be present in the residue.

                    parents = [atom for atom in residue.atoms() if atom.element != elem.hydrogen]
                    parentNames = [atom.name for atom in parents]
                    hydrogens = [h for h in spec.hydrogens if (variant is None and pH <= h.maxph) or (h.variants is None and pH <= h.maxph) or (h.variants is not None and variant in h.variants)]
                    hydrogens = [h for h in hydrogens if h.terminal is None or (isNTerminal and h.terminal == 'N') or (isCTerminal and h.terminal == 'C')]
                    hydrogens = [h for h in hydrogens if h.parent in parentNames]

                    # Loop over atoms in the residue, adding them to the new topology along with required hydrogens.

                    for parent in residue.atoms():
                        # Add the atom.

                        newAtom = newTopology.addAtom(parent.name, parent.element, newResidue)
                        newAtoms[parent] = newAtom
                        newPositions.append(deepcopy(self.positions[parent.index]))
                        if parent in parents:
                            # Match expected hydrogens with existing ones and find which ones need to be added.

                            existing = [atom for atom in bonded[parent] if atom.element == elem.hydrogen]
                            expected = [h for h in hydrogens if h.parent == parent.name]
                            if len(existing) < len(expected):
                                # Try to match up existing hydrogens to expected ones.

                                matches = []
                                for e in existing:
                                    match = [h for h in expected if h.name == e.name]
                                    if len(match) > 0:
                                        matches.append(match[0])
                                        expected.remove(match[0])
                                    else:
                                        matches.append(None)

                                # If any hydrogens couldn't be matched by name, just match them arbitrarily.

                                for i in range(len(matches)):
                                    if matches[i] is None:
                                        matches[i] = expected[-1]
                                        expected.remove(expected[-1])

                                # Add the missing hydrogens.

                                for h in expected:
                                    newH = newTopology.addAtom(h.name, elem.hydrogen, newResidue)
                                    newIndices.append(newH.index)
                                    delta = Vec3(0, 0, 0)*nanometer
                                    if len(bonded[parent]) > 0:
                                        for other in bonded[parent]:
                                            delta += self.positions[parent.index]-self.positions[other.index]
                                    else:
                                        delta = Vec3(random.random(), random.random(), random.random())*nanometer
                                    delta *= 0.1*nanometer/norm(delta)
                                    delta += 0.05*Vec3(random.random(), random.random(), random.random())*nanometer
                                    delta *= 0.1*nanometer/norm(delta)
                                    newPositions.append(self.positions[parent.index]+delta)
                                    newTopology.addBond(newAtom, newH)
                else:
                    # Just copy over the residue.

                    for atom in residue.atoms():
                        newAtom = newTopology.addAtom(atom.name, atom.element, newResidue)
                        newAtoms[atom] = newAtom
                        newPositions.append(deepcopy(self.positions[atom.index]))
        for bond in self.topology.bonds():
            if bond[0] in newAtoms and bond[1] in newAtoms:
                newTopology.addBond(newAtoms[bond[0]], newAtoms[bond[1]])

        # The hydrogens were added at random positions.  Now use the ForceField to fix them up.

        system = forcefield.createSystem(newTopology, rigidWater=False)
        atoms = list(newTopology.atoms())
        for i in range(system.getNumParticles()):
            if atoms[i].element != elem.hydrogen:
                # This is a heavy atom, so make it immobile.
                system.setParticleMass(i, 0)
        if platform is None:
            context = Context(system, VerletIntegrator(0.0))
        else:
            context = Context(system, VerletIntegrator(0.0), platform)
        context.setPositions(newPositions)
        LocalEnergyMinimizer.minimize(context)
        self.topology = newTopology
        self.positions = context.getState(getPositions=True).getPositions()
        return actualVariants
Пример #10
0
class SingleContext:
    """Mimics the MultiContext API but does not spawn worker processes.

    Parameters:
    -----------
    n_workers : int
        Needs to be 1.
    system : simtk.openmm.System
        The system that contains all forces.
    integrator : simtk.openmm.Integrator
        An OpenMM integrator.
    platform_name : str
        The name of an OpenMM platform ('Reference', 'CPU', 'CUDA', or 'OpenCL')
    platform_properties : dict, optional
        A dictionary of platform properties.
    """

    def __init__(self, n_workers, system, integrator, platform_name, platform_properties={}):
        """Set up workers and queues."""
        from simtk.openmm import Platform, Context
        assert n_workers == 1
        openmm_platform = Platform.getPlatformByName(platform_name)
        self._openmm_context = Context(system, integrator, openmm_platform, platform_properties)

    def evaluate(
            self,
            positions,
            box_vectors=None,
            evaluate_energy=True,
            evaluate_force=True,
            evaluate_positions=False,
            evaluate_path_probability_ratio=False,
            err_handling="warning",
            n_simulation_steps=0
    ):
        """Compute energies and/or forces.

        Parameters:
        -----------
        positions : numpy.ndarray
            The particle positions in nanometer; its shape is (batch_size, num_particles, 3).
        box_vectors : numpy.ndarray, optional
            The periodic box vectors in nanometer; its shape is (batch_size, 3, 3).
            If not specified, don't change the box vectors.
        evaluate_energy : bool, optional
            Whether to compute energies.
        evaluate_force : bool, optional
            Whether to compute forces.
        evaluate_positions : bool, optional
            Whether to return positions.
        evaluate_path_probability_ratio : bool, optional
            Whether to compute the log path probability ratio. Makes only sense for PathProbabilityIntegrator instances.
        _err_handling : str, optional
            How to handle infinite energies (one of {"warning", "ignore", "exception"}).
        n_simulation_steps : int, optional
            If > 0, perform a number of simulation steps and compute energy and forces for the resulting state.

        Returns:
        --------
        energies : np.ndarray or None
            The energies in units of kilojoule/mole; its shape  is (len(positions), )
        forces : np.ndarray or None
            The forces in units of kilojoule/mole/nm; its shape is (len(positions), num_particles, 3)
        new_positions : np.ndarray or None
            The positions in units of nm; its shape is (len(positions), num_particles, 3)
        log_path_probability_ratio : np.ndarray or None
            The logarithmic path probability ratios; its shape  is (len(positions), )
        """
        from simtk import unit
        assert box_vectors is None or len(box_vectors) == len(positions), \
            "box_vectors and positions have to be the same length"
        box_vectors = [None for _ in positions] if box_vectors is None else box_vectors

        forces = np.zeros_like(positions)
        energies = np.zeros_like(positions[:,0,0])
        new_positions = np.zeros_like(positions)
        log_path_probability_ratios = np.zeros_like(positions[:,0,0])

        for i, (p, bv) in enumerate(zip(positions, box_vectors)):

            try:
                # initialize state
                self._openmm_context.setPositions(p)
                if bv is not None:
                    self._openmm_context.setPeriodicBoxVectors(bv)
                log_path_probability_ratio = self._openmm_context.getIntegrator().step(n_simulation_steps)

                # compute energy and forces
                state = self._openmm_context.getState(
                    getEnergy=evaluate_energy,
                    getForces=evaluate_force,
                    getPositions=evaluate_positions
                )
                energy = state.getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole) if evaluate_energy else None
                force = (
                    state.getForces(asNumpy=True).value_in_unit(unit.kilojoule_per_mole / unit.nanometer)
                    if evaluate_force else None
                )
                new_pos = state.getPositions().value_in_unit(unit.nanometers) if evaluate_positions else None

                energies[i] = energy if evaluate_energy else 0.0
                forces[i,:,:] = force if evaluate_force else 0.0
                new_positions[i,:,:] = new_pos if evaluate_positions else 0.0
                log_path_probability_ratios[i] = log_path_probability_ratio if evaluate_path_probability_ratio else 0.0

            except Exception as e:
                if err_handling == "warning":
                    warnings.warn("Suppressed exception: {}".format(e))
                elif err_handling == "exception":
                    raise e

        return (
            energies if evaluate_energy else None,
            forces if evaluate_force else None,
            new_positions if evaluate_positions else None,
            log_path_probability_ratios if evaluate_path_probability_ratio else None
        )
Пример #11
0
    class Worker(mp.Process):
        """A worker process that computes energies in its own context.

        Parameters:
        -----------
        task_queue : multiprocessing.Queue
            The queue that the MultiContext pushes tasks to.
        result_queue : multiprocessing.Queue
            The queue that the MultiContext receives results from.
        system : simtk.openmm.System
            The system that contains all forces.
        integrator : simtk.openmm.Integrator
            An OpenMM integrator.
        platform_name : str
            The name of an OpenMM platform ('Reference', 'CPU', 'CUDA', or 'OpenCL')
        platform_properties : dict
            A dictionary of platform properties.
        """

        def __init__(self, task_queue, result_queue, system, integrator, platform_name, platform_properties):
            super(MultiContext.Worker, self).__init__()
            self._task_queue = task_queue
            self._result_queue = result_queue
            self._openmm_system = system
            self._openmm_integrator = pickle.loads( pickle.dumps(integrator))
            self._openmm_platform_name = platform_name
            self._openmm_platform_properties = platform_properties
            self._openmm_context = None

        def run(self):
            """Run the process: set positions and compute energies and forces.
            Positions and box vectors are received from the task_queue in units of nanometers.
            Energies and forces are pushed to the result_queue in units of kJ/mole and kJ/mole/nm, respectively.
            """
            from simtk import unit
            from simtk.openmm import Platform, Context

            # create the context
            # it is crucial to do that in the run function and not in the constructor
            # for some reason, the CPU platform hangs if the context is created in the constructor
            # see also https://github.com/openmm/openmm/issues/2602
            openmm_platform = Platform.getPlatformByName(self._openmm_platform_name)
            self._openmm_context = Context(
                self._openmm_system,
                self._openmm_integrator,
                openmm_platform,
                self._openmm_platform_properties
            )
            self._openmm_context.reinitialize(preserveState=True)

            # get tasks from the task queue
            for task in iter(self._task_queue.get, None):
                (index, positions, box_vectors, evaluate_energy, evaluate_force,
                 evaluate_positions, evaluate_path_probability_ratio, err_handling, n_simulation_steps) = task
                try:
                    # initialize state
                    self._openmm_context.setPositions(positions)
                    if box_vectors is not None:
                        self._openmm_context.setPeriodicBoxVectors(box_vectors)
                    log_path_probability_ratio = self._openmm_integrator.step(n_simulation_steps)

                    # compute energy and forces
                    state = self._openmm_context.getState(
                        getEnergy=evaluate_energy,
                        getForces=evaluate_force,
                        getPositions=evaluate_positions
                    )
                    energy = state.getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole) if evaluate_energy else None
                    forces = (
                        state.getForces(asNumpy=True).value_in_unit(unit.kilojoule_per_mole / unit.nanometer)
                        if evaluate_force else None
                    )
                    new_positions = state.getPositions().value_in_unit(unit.nanometers) if evaluate_positions else None
                except Exception as e:
                    if err_handling == "warning":
                        warnings.warn("Suppressed exception: {}".format(e))
                    elif err_handling == "exception":
                        raise e

                # push energies and forces to the results queue
                self._result_queue.put(
                    [index, energy, forces, new_positions, log_path_probability_ratio]
                )
Пример #12
0
    def addHydrogens(self,
                     forcefield=None,
                     pH=None,
                     variants=None,
                     platform=None):
        """Add missing hydrogens to the model.

        This function automatically changes compatible residues into their constant-pH variant if no variant is specified.:

        Aspartic acid:
            AS4: Form with a 2 hydrogens on each one of the delta oxygens (syn,anti)
                It has 5 titration states.

            Alternative:
            AS2: Has 2 hydrogens (syn, anti) on one of the delta oxygens
                It has 3 titration states.

        Cysteine:
            CYS: Neutral form with a hydrogen on the sulfur
            CYX: No hydrogen on the sulfur (either negatively charged, or part of a disulfide bond)

        Glutamic acid:
            GL4: Form with a 2 hydrogens on each one of the epsilon oxygens (syn,anti)
                It has 5 titration states.

        Histidine:
            HIP: Positively charged form with hydrogens on both ND1 and NE2
                It has 3 titration states.

        The variant to use for each residue is determined by the following rules:

        1. Any Cysteine that participates in a disulfide bond uses the CYX variant regardless of pH.
        2. Other residues are all set to maximally protonated state, which can be updated using a proton drive

        You can override these rules by explicitly specifying a variant for any residue.  To do that, provide a list for the
        'variants' parameter, and set the corresponding element to the name of the variant to use.

        A special case is when the model already contains a hydrogen that should not be present in the desired variant.
        If you explicitly specify a variant using the 'variants' parameter, the residue will be modified to match the
        desired variant, removing hydrogens if necessary.  On the other hand, for residues whose variant is selected
        automatically, this function will only add hydrogens.  It will never remove ones that are already present in the
        model.

        Definitions for standard amino acids and nucleotides are built in.  You can call loadHydrogenDefinitions() to load
        additional definitions for other residue types.

        Parameters
        ----------
        forcefield : ForceField=None
            the ForceField to use for determining the positions of hydrogens.
            If this is None, positions will be picked which are generally
            reasonable but not optimized for any particular ForceField.
        pH : None,
            Kept for compatibility reasons. Has no effect.
        variants : list=None
            an optional list of variants to use.  If this is specified, its
            length must equal the number of residues in the model.  variants[i]
            is the name of the variant to use for residue i (indexed starting at
            0). If an element is None, the standard rules will be followed to
            select a variant for that residue.
        platform : Platform=None
            the Platform to use when computing the hydrogen atom positions.  If
            this is None, the default Platform will be used.

        Returns
        -------
        list
             a list of what variant was actually selected for each residue,
             in the same format as the variants parameter

        Notes
        -----

        This function does not use a pH specification. The argument is kept for compatibility reasons.

        """
        # Check the list of variants.

        if pH is not None:
            print("Ignored pH argument provided for constant-pH residues.")

        residues = list(self.topology.residues())
        if variants is not None:
            if len(variants) != len(residues):
                raise ValueError(
                    "The length of the variants list must equal the number of residues"
                )
        else:
            variants = [None] * len(residues)
        actualVariants = [None] * len(residues)

        # Load the residue specifications.

        if not Modeller._hasLoadedStandardHydrogens:
            Modeller.loadHydrogenDefinitions(
                os.path.join(os.path.dirname(__file__), "data",
                             "hydrogens-amber10-constph.xml"))

        # Make a list of atoms bonded to each atom.

        bonded = {}
        for atom in self.topology.atoms():
            bonded[atom] = []
        for atom1, atom2 in self.topology.bonds():
            bonded[atom1].append(atom2)
            bonded[atom2].append(atom1)

        # Define a function that decides whether a set of atoms form a hydrogen bond, using fairly tolerant criteria.

        def isHbond(d, h, a):
            if norm(d - a) > 0.35 * nanometer:
                return False
            deltaDH = h - d
            deltaHA = a - h
            deltaDH /= norm(deltaDH)
            deltaHA /= norm(deltaHA)
            return acos(dot(deltaDH, deltaHA)) < 50 * degree

        # Loop over residues.

        newTopology = Topology()
        newTopology.setPeriodicBoxVectors(
            self.topology.getPeriodicBoxVectors())
        newAtoms = {}
        newPositions = [] * nanometer
        newIndices = []
        acceptors = [
            atom for atom in self.topology.atoms()
            if atom.element in (elem.oxygen, elem.nitrogen)
        ]
        for chain in self.topology.chains():
            newChain = newTopology.addChain(chain.id)
            for residue in chain.residues():
                newResidue = newTopology.addResidue(residue.name, newChain,
                                                    residue.id)
                isNTerminal = residue == chain._residues[0]
                isCTerminal = residue == chain._residues[-1]
                if residue.name in Modeller._residueHydrogens:
                    # Add hydrogens.  First select which variant to use.

                    spec = Modeller._residueHydrogens[residue.name]
                    variant = variants[residue.index]
                    if variant is None:
                        if residue.name == "CYS":
                            # If this is part of a disulfide, use CYX.

                            sulfur = [
                                atom for atom in residue.atoms()
                                if atom.element == elem.sulfur
                            ]
                            if len(sulfur) == 1 and any(
                                (atom.residue != residue
                                 for atom in bonded[sulfur[0]])):
                                variant = "CYX"
                        if residue.name == "HIS":
                            variant = "HIP"
                        if residue.name == "GLU":
                            variant = "GL4"
                        if residue.name == "ASP":
                            variant = "AS4"
                    if variant is not None and variant not in spec.variants:
                        raise ValueError("Illegal variant for %s residue: %s" %
                                         (residue.name, variant))
                    actualVariants[residue.index] = variant
                    removeExtraHydrogens = variants[residue.index] is not None

                    # Make a list of hydrogens that should be present in the residue.

                    parents = [
                        atom for atom in residue.atoms()
                        if atom.element != elem.hydrogen
                    ]
                    parentNames = [atom.name for atom in parents]
                    hydrogens = [
                        h for h in spec.hydrogens
                        if (variant is None) or (h.variants is None) or (
                            h.variants is not None and variant in h.variants)
                    ]
                    hydrogens = [
                        h for h in hydrogens if h.terminal is None or (
                            isNTerminal and h.terminal == "N") or (
                                isCTerminal and h.terminal == "C")
                    ]
                    hydrogens = [
                        h for h in hydrogens if h.parent in parentNames
                    ]

                    # Loop over atoms in the residue, adding them to the new topology along with required hydrogens.

                    for parent in residue.atoms():
                        # Check whether this is a hydrogen that should be removed.

                        if (removeExtraHydrogens
                                and parent.element == elem.hydrogen
                                and not any(parent.name == h.name
                                            for h in hydrogens)):
                            continue

                        # Add the atom.

                        newAtom = newTopology.addAtom(parent.name,
                                                      parent.element,
                                                      newResidue)
                        newAtoms[parent] = newAtom
                        newPositions.append(
                            deepcopy(self.positions[parent.index]))
                        if parent in parents:
                            # Match expected hydrogens with existing ones and find which ones need to be added.

                            existing = [
                                atom for atom in bonded[parent]
                                if atom.element == elem.hydrogen
                            ]
                            expected = [
                                h for h in hydrogens if h.parent == parent.name
                            ]
                            if len(existing) < len(expected):
                                # Try to match up existing hydrogens to expected ones.

                                matches = []
                                for e in existing:
                                    match = [
                                        h for h in expected if h.name == e.name
                                    ]
                                    if len(match) > 0:
                                        matches.append(match[0])
                                        expected.remove(match[0])
                                    else:
                                        matches.append(None)

                                # If any hydrogens couldn't be matched by name, just match them arbitrarily.

                                for i in range(len(matches)):
                                    if matches[i] is None:
                                        matches[i] = expected[-1]
                                        expected.remove(expected[-1])

                                # Add the missing hydrogens.

                                for h in expected:
                                    newH = newTopology.addAtom(
                                        h.name, elem.hydrogen, newResidue)
                                    newIndices.append(newH.index)
                                    delta = Vec3(0, 0, 0) * nanometer
                                    if len(bonded[parent]) > 0:
                                        for other in bonded[parent]:
                                            delta += (
                                                self.positions[parent.index] -
                                                self.positions[other.index])
                                    else:
                                        delta = (Vec3(
                                            random.random(),
                                            random.random(),
                                            random.random(),
                                        ) * nanometer)
                                    delta *= 0.1 * nanometer / norm(delta)
                                    delta += (0.05 * Vec3(
                                        random.random(),
                                        random.random(),
                                        random.random(),
                                    ) * nanometer)
                                    delta *= 0.1 * nanometer / norm(delta)
                                    newPositions.append(
                                        self.positions[parent.index] + delta)
                                    newTopology.addBond(newAtom, newH)
                else:
                    # Just copy over the residue.

                    for atom in residue.atoms():
                        newAtom = newTopology.addAtom(atom.name, atom.element,
                                                      newResidue)
                        newAtoms[atom] = newAtom
                        newPositions.append(
                            deepcopy(self.positions[atom.index]))
        for bond in self.topology.bonds():
            if bond[0] in newAtoms and bond[1] in newAtoms:
                newTopology.addBond(newAtoms[bond[0]], newAtoms[bond[1]])

        # The hydrogens were added at random positions.  Now perform an energy minimization to fix them up.

        if forcefield is not None:
            # Use the ForceField the user specified.

            system = forcefield.createSystem(newTopology, rigidWater=False)
            atoms = list(newTopology.atoms())
            for i in range(system.getNumParticles()):
                if atoms[i].element != elem.hydrogen:
                    # This is a heavy atom, so make it immobile.
                    system.setParticleMass(i, 0)
        else:
            # Create a System that restrains the distance of each hydrogen from its parent atom
            # and causes hydrogens to spread out evenly.

            system = System()
            nonbonded = CustomNonbondedForce("100/((r/0.1)^4+1)")
            bonds = HarmonicBondForce()
            angles = HarmonicAngleForce()
            system.addForce(nonbonded)
            system.addForce(bonds)
            system.addForce(angles)
            bondedTo = []
            for atom in newTopology.atoms():
                nonbonded.addParticle([])
                if atom.element != elem.hydrogen:
                    system.addParticle(0.0)
                else:
                    system.addParticle(1.0)
                bondedTo.append([])
            for atom1, atom2 in newTopology.bonds():
                if atom1.element == elem.hydrogen or atom2.element == elem.hydrogen:
                    bonds.addBond(atom1.index, atom2.index, 0.1, 100_000.0)
                bondedTo[atom1.index].append(atom2)
                bondedTo[atom2.index].append(atom1)
            for residue in newTopology.residues():
                if residue.name == "HOH":
                    # Add an angle term to make the water geometry correct.

                    atoms = list(residue.atoms())
                    oindex = [
                        i for i in range(len(atoms))
                        if atoms[i].element == elem.oxygen
                    ]
                    if len(atoms) == 3 and len(oindex) == 1:
                        hindex = list(set([0, 1, 2]) - set(oindex))
                        angles.addAngle(
                            atoms[hindex[0]].index,
                            atoms[oindex[0]].index,
                            atoms[hindex[1]].index,
                            1.824,
                            836.8,
                        )
                else:
                    # Add angle terms for any hydroxyls.

                    for atom in residue.atoms():
                        index = atom.index
                        if (atom.element == elem.oxygen
                                and len(bondedTo[index]) == 2 and elem.hydrogen
                                in (a.element for a in bondedTo[index])):
                            angles.addAngle(
                                bondedTo[index][0].index,
                                index,
                                bondedTo[index][1].index,
                                1.894,
                                460.24,
                            )

        if platform is None:
            context = Context(system, VerletIntegrator(0.0))
        else:
            context = Context(system, VerletIntegrator(0.0), platform)
        context.setPositions(newPositions)
        LocalEnergyMinimizer.minimize(context, 1.0, 50)
        self.topology = newTopology
        self.positions = context.getState(getPositions=True).getPositions()
        del context
        return actualVariants
Пример #13
0
    def addHydrogens(self, forcefield=None, pH=None, variants=None, platform=None):
        """Add missing hydrogens to the model.

        This function automatically changes compatible residues into their constant-pH variant if no variant is specified.:

        Aspartic acid:
            AS4: Form with a 2 hydrogens on each one of the delta oxygens (syn,anti)
                It has 5 titration states.

            Alternative:
            AS2: Has 2 hydrogens (syn, anti) on one of the delta oxygens
                It has 3 titration states.

        Cysteine:
            CYS: Neutral form with a hydrogen on the sulfur
            CYX: No hydrogen on the sulfur (either negatively charged, or part of a disulfide bond)

        Glutamic acid:
            GL4: Form with a 2 hydrogens on each one of the epsilon oxygens (syn,anti)
                It has 5 titration states.

        Histidine:
            HIP: Positively charged form with hydrogens on both ND1 and NE2
                It has 3 titration states.

        The variant to use for each residue is determined by the following rules:

        1. Any Cysteine that participates in a disulfide bond uses the CYX variant regardless of pH.
        2. Other residues are all set to maximally protonated state, which can be updated using a proton drive

        You can override these rules by explicitly specifying a variant for any residue.  To do that, provide a list for the
        'variants' parameter, and set the corresponding element to the name of the variant to use.

        A special case is when the model already contains a hydrogen that should not be present in the desired variant.
        If you explicitly specify a variant using the 'variants' parameter, the residue will be modified to match the
        desired variant, removing hydrogens if necessary.  On the other hand, for residues whose variant is selected
        automatically, this function will only add hydrogens.  It will never remove ones that are already present in the
        model.

        Definitions for standard amino acids and nucleotides are built in.  You can call loadHydrogenDefinitions() to load
        additional definitions for other residue types.

        Parameters
        ----------
        forcefield : ForceField=None
            the ForceField to use for determining the positions of hydrogens.
            If this is None, positions will be picked which are generally
            reasonable but not optimized for any particular ForceField.
        pH : None,
            Kept for compatibility reasons. Has no effect.
        variants : list=None
            an optional list of variants to use.  If this is specified, its
            length must equal the number of residues in the model.  variants[i]
            is the name of the variant to use for residue i (indexed starting at
            0). If an element is None, the standard rules will be followed to
            select a variant for that residue.
        platform : Platform=None
            the Platform to use when computing the hydrogen atom positions.  If
            this is None, the default Platform will be used.

        Returns
        -------
        list
             a list of what variant was actually selected for each residue,
             in the same format as the variants parameter

        Notes
        -----

        This function does not use a pH specification. The argument is kept for compatibility reasons.

        """
        # Check the list of variants.

        if pH is not None:
            print("Ignored pH argument provided for constant-pH residues.")

        residues = list(self.topology.residues())
        if variants is not None:
            if len(variants) != len(residues):
                raise ValueError(
                    "The length of the variants list must equal the number of residues"
                )
        else:
            variants = [None] * len(residues)
        actualVariants = [None] * len(residues)

        # Load the residue specifications.

        if not Modeller._hasLoadedStandardHydrogens:
            Modeller.loadHydrogenDefinitions(
                os.path.join(
                    os.path.dirname(__file__), "data", "hydrogens-amber10-constph.xml"
                )
            )

        # Make a list of atoms bonded to each atom.

        bonded = {}
        for atom in self.topology.atoms():
            bonded[atom] = []
        for atom1, atom2 in self.topology.bonds():
            bonded[atom1].append(atom2)
            bonded[atom2].append(atom1)

        # Define a function that decides whether a set of atoms form a hydrogen bond, using fairly tolerant criteria.

        def isHbond(d, h, a):
            if norm(d - a) > 0.35 * nanometer:
                return False
            deltaDH = h - d
            deltaHA = a - h
            deltaDH /= norm(deltaDH)
            deltaHA /= norm(deltaHA)
            return acos(dot(deltaDH, deltaHA)) < 50 * degree

        # Loop over residues.

        newTopology = Topology()
        newTopology.setPeriodicBoxVectors(self.topology.getPeriodicBoxVectors())
        newAtoms = {}
        newPositions = [] * nanometer
        newIndices = []
        acceptors = [
            atom
            for atom in self.topology.atoms()
            if atom.element in (elem.oxygen, elem.nitrogen)
        ]
        for chain in self.topology.chains():
            newChain = newTopology.addChain(chain.id)
            for residue in chain.residues():
                newResidue = newTopology.addResidue(residue.name, newChain, residue.id)
                isNTerminal = residue == chain._residues[0]
                isCTerminal = residue == chain._residues[-1]
                if residue.name in Modeller._residueHydrogens:
                    # Add hydrogens.  First select which variant to use.

                    spec = Modeller._residueHydrogens[residue.name]
                    variant = variants[residue.index]
                    if variant is None:
                        if residue.name == "CYS":
                            # If this is part of a disulfide, use CYX.

                            sulfur = [
                                atom
                                for atom in residue.atoms()
                                if atom.element == elem.sulfur
                            ]
                            if len(sulfur) == 1 and any(
                                (atom.residue != residue for atom in bonded[sulfur[0]])
                            ):
                                variant = "CYX"
                        if residue.name == "HIS":
                            variant = "HIP"
                        if residue.name == "GLU":
                            variant = "GL4"
                        if residue.name == "ASP":
                            variant = "AS4"
                    if variant is not None and variant not in spec.variants:
                        raise ValueError(
                            "Illegal variant for %s residue: %s"
                            % (residue.name, variant)
                        )
                    actualVariants[residue.index] = variant
                    removeExtraHydrogens = variants[residue.index] is not None

                    # Make a list of hydrogens that should be present in the residue.

                    parents = [
                        atom
                        for atom in residue.atoms()
                        if atom.element != elem.hydrogen
                    ]
                    parentNames = [atom.name for atom in parents]
                    hydrogens = [
                        h
                        for h in spec.hydrogens
                        if (variant is None)
                        or (h.variants is None)
                        or (h.variants is not None and variant in h.variants)
                    ]
                    hydrogens = [
                        h
                        for h in hydrogens
                        if h.terminal is None
                        or (isNTerminal and h.terminal == "N")
                        or (isCTerminal and h.terminal == "C")
                    ]
                    hydrogens = [h for h in hydrogens if h.parent in parentNames]

                    # Loop over atoms in the residue, adding them to the new topology along with required hydrogens.

                    for parent in residue.atoms():
                        # Check whether this is a hydrogen that should be removed.

                        if (
                            removeExtraHydrogens
                            and parent.element == elem.hydrogen
                            and not any(parent.name == h.name for h in hydrogens)
                        ):
                            continue

                        # Add the atom.

                        newAtom = newTopology.addAtom(
                            parent.name, parent.element, newResidue
                        )
                        newAtoms[parent] = newAtom
                        newPositions.append(deepcopy(self.positions[parent.index]))
                        if parent in parents:
                            # Match expected hydrogens with existing ones and find which ones need to be added.

                            existing = [
                                atom
                                for atom in bonded[parent]
                                if atom.element == elem.hydrogen
                            ]
                            expected = [h for h in hydrogens if h.parent == parent.name]
                            if len(existing) < len(expected):
                                # Try to match up existing hydrogens to expected ones.

                                matches = []
                                for e in existing:
                                    match = [h for h in expected if h.name == e.name]
                                    if len(match) > 0:
                                        matches.append(match[0])
                                        expected.remove(match[0])
                                    else:
                                        matches.append(None)

                                # If any hydrogens couldn't be matched by name, just match them arbitrarily.

                                for i in range(len(matches)):
                                    if matches[i] is None:
                                        matches[i] = expected[-1]
                                        expected.remove(expected[-1])

                                # Add the missing hydrogens.

                                for h in expected:
                                    newH = newTopology.addAtom(
                                        h.name, elem.hydrogen, newResidue
                                    )
                                    newIndices.append(newH.index)
                                    delta = Vec3(0, 0, 0) * nanometer
                                    if len(bonded[parent]) > 0:
                                        for other in bonded[parent]:
                                            delta += (
                                                self.positions[parent.index]
                                                - self.positions[other.index]
                                            )
                                    else:
                                        delta = (
                                            Vec3(
                                                random.random(),
                                                random.random(),
                                                random.random(),
                                            )
                                            * nanometer
                                        )
                                    delta *= 0.1 * nanometer / norm(delta)
                                    delta += (
                                        0.05
                                        * Vec3(
                                            random.random(),
                                            random.random(),
                                            random.random(),
                                        )
                                        * nanometer
                                    )
                                    delta *= 0.1 * nanometer / norm(delta)
                                    newPositions.append(
                                        self.positions[parent.index] + delta
                                    )
                                    newTopology.addBond(newAtom, newH)
                else:
                    # Just copy over the residue.

                    for atom in residue.atoms():
                        newAtom = newTopology.addAtom(
                            atom.name, atom.element, newResidue
                        )
                        newAtoms[atom] = newAtom
                        newPositions.append(deepcopy(self.positions[atom.index]))
        for bond in self.topology.bonds():
            if bond[0] in newAtoms and bond[1] in newAtoms:
                newTopology.addBond(newAtoms[bond[0]], newAtoms[bond[1]])

        # The hydrogens were added at random positions.  Now perform an energy minimization to fix them up.

        if forcefield is not None:
            # Use the ForceField the user specified.

            system = forcefield.createSystem(newTopology, rigidWater=False)
            atoms = list(newTopology.atoms())
            for i in range(system.getNumParticles()):
                if atoms[i].element != elem.hydrogen:
                    # This is a heavy atom, so make it immobile.
                    system.setParticleMass(i, 0)
        else:
            # Create a System that restrains the distance of each hydrogen from its parent atom
            # and causes hydrogens to spread out evenly.

            system = System()
            nonbonded = CustomNonbondedForce("100/((r/0.1)^4+1)")
            bonds = HarmonicBondForce()
            angles = HarmonicAngleForce()
            system.addForce(nonbonded)
            system.addForce(bonds)
            system.addForce(angles)
            bondedTo = []
            for atom in newTopology.atoms():
                nonbonded.addParticle([])
                if atom.element != elem.hydrogen:
                    system.addParticle(0.0)
                else:
                    system.addParticle(1.0)
                bondedTo.append([])
            for atom1, atom2 in newTopology.bonds():
                if atom1.element == elem.hydrogen or atom2.element == elem.hydrogen:
                    bonds.addBond(atom1.index, atom2.index, 0.1, 100_000.0)
                bondedTo[atom1.index].append(atom2)
                bondedTo[atom2.index].append(atom1)
            for residue in newTopology.residues():
                if residue.name == "HOH":
                    # Add an angle term to make the water geometry correct.

                    atoms = list(residue.atoms())
                    oindex = [
                        i for i in range(len(atoms)) if atoms[i].element == elem.oxygen
                    ]
                    if len(atoms) == 3 and len(oindex) == 1:
                        hindex = list(set([0, 1, 2]) - set(oindex))
                        angles.addAngle(
                            atoms[hindex[0]].index,
                            atoms[oindex[0]].index,
                            atoms[hindex[1]].index,
                            1.824,
                            836.8,
                        )
                else:
                    # Add angle terms for any hydroxyls.

                    for atom in residue.atoms():
                        index = atom.index
                        if (
                            atom.element == elem.oxygen
                            and len(bondedTo[index]) == 2
                            and elem.hydrogen in (a.element for a in bondedTo[index])
                        ):
                            angles.addAngle(
                                bondedTo[index][0].index,
                                index,
                                bondedTo[index][1].index,
                                1.894,
                                460.24,
                            )

        if platform is None:
            context = Context(system, VerletIntegrator(0.0))
        else:
            context = Context(system, VerletIntegrator(0.0), platform)
        context.setPositions(newPositions)
        LocalEnergyMinimizer.minimize(context, 1.0, 50)
        self.topology = newTopology
        self.positions = context.getState(getPositions=True).getPositions()
        del context
        return actualVariants