Add dihedral MCMove class

.github/workflows/CI.yml

@@ -84,7 +84,7 @@ jobs:
         # conda setup requires this special shell
         shell: bash -l {0}
         run: |
-          python -m pip install . --no-deps
+          python -m pip install . --no-deps -vv
           conda list
 
       - name: Run tests

docs/mcmc.rst

@@ -136,3 +136,4 @@ A number of MCMC component move types that can be arranged into groups or subcla
     MonteCarloBarostatMove
     MCDisplacementMove
     MCRotationMove
+    MCDihedralRotationMove

docs/releasehistory.rst

@@ -1,6 +1,12 @@
 Release History
 ***************
 
+0.21.0 - MC Dihedral Rotation Move
+========================================
+
+Enhancements
+------------
+- Add `MCDihedralRotationMove`, an extension of MetropolizedMove that allows rigid rotation moves about a specified dihedral
 
 0.20.0 - Periodic alchemical integrators
 ========================================

openmmtools/mcmc.py

@@ -116,6 +116,8 @@
 
 import numpy as np
 from simtk import openmm, unit
+import math
+import random
 
 from openmmtools import integrators, cache, utils
 from openmmtools.utils import SubhookedABCMeta, Timer
@@ -1894,6 +1896,213 @@ def _propose_positions(self, initial_positions):
         """Implement MetropolizedMove._propose_positions for apply()"""
         return self.rotate_positions(initial_positions)
 
+# =============================================================================
+# RANDOM DIHEDRAL ROTATION MOVE
+# =============================================================================
+
+class MCDihedralRotationMove(MetropolizedMove):
+    """A metropolized move that randomly rotates a subset of atoms about a specified bond.
+
+    Parameters
+    ----------
+    atom_subset : slice or list of int
+        Atom indices where the first four are the dihedral of interest. The rotatable bond is specified by
+        the 2nd and 3rd atom indices in the list. All atom indices after the 3rd correspond to the atoms to be rotated.
+        Example: [6, 8, 10, 18, 13, 14, 15, 16, 17, 19] means that the dihedral of interest is 6-8-10-18,
+        the rotatable bond is 8-10, and the atoms to be rotated are 18, 13, 14, 15, 16, 17, 19.
+    desired_angle : float, default -np.inf
+        Desired angle to rotate to. If not specified, this will be chosen randomly. Must be between -pi and pi.
+        Example: If the current angle is pi, and desired_angle is 0, the atoms will be rotated by -pi.
+    context_cache : openmmtools.cache.ContextCache, optional
+        The ContextCache to use for Context creation. If None, the global cache
+        openmmtools.cache.global_context_cache is used (default is None).
+
+    Attributes
+    ----------
+    atom_subset
+    desired_angle
+    context_cache
+    n_accepted : int
+        The number of proposals accepted.
+    n_proposed : int
+        The total number of attempted moves.
+    proposed_positions : np.array of floats
+        Proposed positions of atoms in atom_subset
+
+    See Also
+    --------
+    MetropolizedMove
+
+    Examples
+    --------
+
+    >>> import numpy as np
+    >>> from simtk import unit
+    >>> from openmmtools import testsystems
+    >>> from openmmtools.states import ThermodynamicState, SamplerState
+
+    Create and run an alanine dipeptide simulation with a sidechain dihedral move.
+
+    >>> test = testsystems.AlanineDipeptideVacuum()
+    >>> thermodynamic_state = ThermodynamicState(system=test.system,
+    ...                                          temperature=300*unit.kelvin)
+    >>> sampler_state = SamplerState(positions=test.positions)
+    >>> # Create a dihedral move
+    >>> move = MCDihedralRotationMove([6, 8, 10, 11, 12, 13], desired_angle=-0.95549) # Rotate the sidechain by 2pi/3
+    >>> # Create an MCMC sampler instance and run 2 iterations of the simulation.
+    >>> sampler = MCMCSampler(thermodynamic_state, sampler_state, move=move)
+    >>> sampler.minimize()
+    >>> sampler.run(n_iterations=2)
+    """
+
+    def __init__(self, atom_subset, desired_angle=-np.inf, context_cache=None):
+        self.atom_subset = atom_subset
+        self.desired_angle = desired_angle
+        self.context_cache = context_cache
+        self.n_accepted = 0
+        self.n_proposed = 0
+        self.proposed_positions = None
+
+    @staticmethod
+    def generate_rotation_matrix(axis, theta):
+        """Generate the rotation matrix associated with counterclockwise rotation
+        about the given axis by theta radians.
+
+        Parameters
+        ----------
+        axis : np.array of three floats
+            The axis of rotation (unitless)
+        theta : float
+            Rotation angle in radians
+
+        Returns
+        -------
+        np.array of floats
+            Rotation matrix
+
+        """
+        axis = np.asarray(axis)
+        axis = axis / math.sqrt(np.dot(axis, axis))
+        a = math.cos(theta / 2.0)
+        b, c, d = -axis * math.sin(theta / 2.0)
+        aa, bb, cc, dd = a * a, b * b, c * c, d * d
+        bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
+        return np.array([[aa + bb - cc - dd, 2 * (bc + ad),
+                          2 * (bd - ac)], [2 * (bc - ad), aa + cc - bb - dd, 2 * (cd + ab)],
+                         [2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc]])
+
+    @staticmethod
+    def compute_dihedral(dihedral):
+        """Given the positions of four atoms that form a dihedral, compute the angle in radians.
+
+        Parameters
+        ----------
+        dihedral : np.array or list of floats
+            Positions of atoms forming the dihedral
+
+        Returns
+        -------
+        float
+            Dihedral angle in radians
+
+        """
+
+        def _cross_vec3(a, b):
+            c = np.zeros(3)
+            c[0] = a[1] * b[2] - a[2] * b[1]
+            c[1] = a[2] * b[0] - a[0] * b[2]
+            c[2] = a[0] * b[1] - a[1] * b[0]
+            return c
+
+        def _norm(a):
+            n_2 = np.dot(a, a)
+            return np.sqrt(n_2)
+
+        atom_position = dihedral[0] / dihedral[0].unit
+        bond_position = dihedral[1] / dihedral[1].unit
+        angle_position = dihedral[2] / dihedral[2].unit
+        torsion_position = dihedral[3] / dihedral[3].unit
+
+        a = atom_position - bond_position
+        b = angle_position - bond_position
+        c = angle_position - torsion_position
+        a_u = a / _norm(a)
+        b_u = b / _norm(b)
+        c_u = c / _norm(c)
+
+        plane1 = _cross_vec3(a_u, b_u)
+        plane2 = _cross_vec3(b_u, c_u)
+
+        cos_phi = np.dot(plane1, plane2) / (_norm(plane1) * _norm(plane2))
+        if cos_phi < -1.0:
+            cos_phi = -1.0
+        elif cos_phi > 1.0:
+            cos_phi = 1.0
+
+        phi = np.arccos(cos_phi)
+
+        if np.dot(a, plane2) <= 0:
+            phi = -phi
+
+        return phi
+
+    def rotate_positions(self, initial_positions):
+        """Apply rotation to atom_subset positions.
+
+        Parameters
+        ----------
+        initial_positions : numpy.ndarray simtk.unit.Quantity
+            The positions of all atoms in atom_subset.
+
+        Returns
+        -------
+        rotated_positions : numpy.ndarray simtk.unit.Quantity
+            The rotated positions.
+        """
+
+        old_angle = self.compute_dihedral(initial_positions[:4])
+
+        if self.desired_angle == -np.inf:
+            # Choose a random rotation angle
+            theta = random.uniform(-math.pi, math.pi)
+        else:
+            # If desired_angle is specified, determine rotation angle
+            if not (self.desired_angle <= math.pi and self.desired_angle >= -math.pi):
+                raise Exception("Desired angle must be less than pi and greater than -pi")
+            theta = self.desired_angle - old_angle
+        logger.info(f"Rotating by {theta} radians")
+
+        # Make a copy of the initial positions
+        new_positions = copy.deepcopy(initial_positions)
+
+        # Find the rotation axis using the initial positions
+        axis1 = 1
+        axis2 = 2
+        rotation_axis = (initial_positions[axis1] - initial_positions[axis2]) / initial_positions.unit
+
+        # Calculate the rotation matrix
+        rotation_matrix = self.generate_rotation_matrix(rotation_axis, theta)
+
+        # Apply the rotation matrix to the target atoms
+        for atom_index in range(3, len(self.atom_subset)):
+            # Find the reduced position (substract out axis)
+            reduced_position = (initial_positions[atom_index] - initial_positions[axis2])._value
+
+            # Find the new positions by multiplying by rot matrix
+            new_position = np.dot(rotation_matrix, reduced_position) * initial_positions.unit + initial_positions[axis2]
+
+            # Update the new positions
+            new_positions[atom_index][0] = new_position[0]
+            new_positions[atom_index][1] = new_position[1]
+            new_positions[atom_index][2] = new_position[2]
+
+        return new_positions
+
+    def _propose_positions(self, initial_positions):
+        """Implement MetropolizedMove._propose_positions for apply()"""
+        self.proposed_positions = self.rotate_positions(initial_positions)
+        return self.proposed_positions
+
 
 # =============================================================================
 # MAIN AND TESTS

openmmtools/tests/test_mcmc.py

@@ -23,7 +23,6 @@
 from openmmtools.states import SamplerState, ThermodynamicState
 from openmmtools.mcmc import *
 
-
 # =============================================================================
 # GLOBAL TEST CONSTANTS
 # =============================================================================
@@ -336,12 +335,13 @@ def test_metropolized_moves():
     original_sampler_state = SamplerState(testsystem.positions)
     thermodynamic_state = ThermodynamicState(testsystem.system, 300*unit.kelvin)
 
-    all_metropolized_moves = MetropolizedMove.__subclasses__()
+    # Test subclasses of MetropolizedMove except MCDihedralRotationMove
+    all_metropolized_moves = [move_class for move_class in MetropolizedMove.__subclasses__() if move_class.__name__ != 'MCDihedralRotationMove']
     for move_class in all_metropolized_moves:
         move = move_class(atom_subset=range(thermodynamic_state.n_particles))
         sampler_state = copy.deepcopy(original_sampler_state)
 
-        # Start applying the move and remove one at each iteration tyring
+        # Start applying the move and remove one at each iteration trying
         # to generate both an accepted and rejected move.
         old_n_accepted, old_n_proposed = 0, 0
         while len(move.atom_subset) > 0:
@@ -364,6 +364,50 @@ def test_metropolized_moves():
         # Check that we were able to generate both an accepted and a rejected move.
         assert len(move.atom_subset) != 0, ('Could not generate an accepted and rejected '
                                             'move for class {}'.format(move_class.__name__))
+    # Test MCDihedralRotationMove
+    # Generate an accepted move
+    move = MCDihedralRotationMove([6, 8, 10, 11, 12, 13], desired_angle=-0.95549) # Rotate the sidechain by 2pi/3, which should always be accepted
+    sampler_state = copy.deepcopy(original_sampler_state)
+    initial_positions = copy.deepcopy(sampler_state.positions)
+    move.apply(thermodynamic_state, sampler_state)
+    final_positions = copy.deepcopy(sampler_state.positions)
+    assert move.n_accepted == 1, ('Rotating alanine dipeptide sidechain by 2pi/3 was not accepted.')
+    if move.n_accepted > old_n_accepted:
+        assert not np.allclose(initial_positions, final_positions)
+
+    # Generate a rejected move
+    for i in range(10):
+        move = MCDihedralRotationMove([6, 8, 10, 11, 12, 13]) # Choose a random rotation
+        move.apply(thermodynamic_state, sampler_state)
+        if move.n_accepted < move.n_proposed:
+            break
+    assert move.n_proposed - move.n_accepted == 1, ('Could not generate a rejected move in 10 iterations.')
+
+def test_generate_rotation_matrix():
+    """Test that MCDihedralRotationMove.generate_rotation_matrix() generates the rotation matrix for rotating the alanine sidechain by 2pi/3 correctly. """
+    theta = -(2*math.pi)/3
+    axis = np.array([-0.0808399, 0.03868687, 0.12475653])
+    rotation_matrix = np.array([[-0.08456292, 0.50454367, -0.85923501],
+                                [-0.90216812, -0.40485611, -0.14894367],
+                                [-0.42301512,  0.76257932, 0.48941903]])
+    assert np.allclose(MCDihedralRotationMove.generate_rotation_matrix(axis, theta), rotation_matrix)
+
+def test_compute_dihedral():
+    """Test that MCDihedralRotationMove.compute_dihedral() computes the N-CA-CB-HB1 dihedral of alanine correctly. """
+    testsystem = testsystems.AlanineDipeptideVacuum()
+    assert np.allclose(MCDihedralRotationMove.compute_dihedral(testsystem.positions[[6, 8, 10, 11]]), 1.04719)
+
+def test_rotate_positions():
+    testsystem = testsystems.AlanineDipeptideVacuum()
+    atom_subset = [6, 8, 10, 11, 12, 13]
+    move = MCDihedralRotationMove(atom_subset, desired_angle=-0.95549) # Rotate the sidechain by 2pi/3
+    new_positions = np.array([[3.55537540e-01, 3.96964880e-01, -3.10000000e-07],
+                            [4.85326210e-01, 4.61392530e-01, -4.30000000e-07],
+                            [5.66130440e-01, 4.22084250e-01, -1.23214800e-01],
+                            [6.59099143e-01, 4.78943429e-01, -1.25418242e-01],
+                            [5.88841910e-01, 3.15546115e-01, -1.19365331e-01],
+                            [5.08287672e-01, 4.43627550e-01, -2.13054053e-01]])
+    assert np.allclose(move.rotate_positions(testsystem.positions[atom_subset]), new_positions)
 
 
 def test_langevin_splitting_move():