Example #1
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)
Example #2
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