Side chain bias - pull request

blues/example.py

@@ -35,7 +35,7 @@ def sidechain_example(yaml_file):
     structure = cfg['Structure']
 
     #Select move type
-    sidechain = SideChainMove(structure, [1])
+    sidechain = SideChainMove(structure, [1], verbose=False)
     #Iniitialize object that selects movestep
     sidechain_mover = MoveEngine(sidechain)
 
@@ -62,5 +62,5 @@ def sidechain_example(yaml_file):
             output.write("%s\n" % str(value)[1:-1])
 
 
-ligrot_example(get_data_filename('blues', '../examples/rotmove_cuda.yml'))
-#sidechain_example(get_data_filename('blues', '../examples/sidechain_cuda.yml'))
+#ligrot_example(get_data_filename('blues', '../examples/rotmove_cuda.yml'))
+sidechain_example(get_data_filename('blues', '../examples/sidechain_cuda.yml'))

blues/simulation.py

@@ -1062,36 +1062,41 @@ def _stepNCMC(self, nstepsNC, moveStep, move_engine=None):
         self._ncmc_sim.currentIter = self.currentIter
         move_engine.selectMove()
 
-        lastStep = nstepsNC - 1
-        for step in range(int(nstepsNC)):
-            try:
-                #Attempt anything related to the move before protocol is performed
-                if not step:
-                    self._ncmc_sim.context = move_engine.selected_move.beforeMove(self._ncmc_sim.context)
-
-                # Attempt selected MoveEngine Move at the halfway point
-                #to ensure protocol is symmetric
-                if step == moveStep:
-                    if hasattr(logger, 'report'):
-                        logger.info = logger.report
-                    #Do move
-                    logger.info('Performing %s...' % move_engine.move_name)
-                    self._ncmc_sim.context = move_engine.runEngine(self._ncmc_sim.context)
-
-                # Do 1 NCMC step with the integrator
-                self._ncmc_sim.step(1)
-
-                #Attempt anything related to the move after protocol is performed
-                if step == lastStep:
-                    self._ncmc_sim.context = move_engine.selected_move.afterMove(self._ncmc_sim.context)
+        #Attempt anything related to the move before protocol is performed
+        self._ncmc_sim.context = move_engine.selected_move.beforeMove(self._ncmc_sim.context)
+        acceptance_ratio = move_engine.selected_move.acceptance_ratio
+        #check if acceptance_ratio is non-zero, if so perform integration
+        if not np.isclose(0, acceptance_ratio):
+
+            for step in range(int(nstepsNC)):
+                try:
+                    # Attempt selected MoveEngine Move at the halfway point
+                    #to ensure protocol is symmetric
+                    if step == moveStep:
+                        if hasattr(logger, 'report'):
+                            logger.info = logger.report
+                        #Do move
+                        logger.info('Performing %s...' % move_engine.move_name)
+                        self._ncmc_sim.context = move_engine.runEngine(self._ncmc_sim.context)
+                        if np.isclose(0, acceptance_ratio):
+                            break
+
+                    # Do 1 NCMC step with the integrator
+                    self._ncmc_sim.step(1)
+
+                except Exception as e:
+                    logger.error(e)
+                    move_engine.selected_move._error(self._ncmc_sim.context)
+                    break
+            #Attempt anything related to the move after protocol is performed
+            self._ncmc_sim.context = move_engine.selected_move.afterMove(self._ncmc_sim.context)
 
-            except Exception as e:
-                logger.error(e)
-                move_engine.selected_move._error(self._ncmc_sim.context)
-                break
+        else:
+            logger.info('move acceptance_ratio is 0; skipping integration')
 
         # ncmc_state1 stores the state AFTER a proposed move.
         ncmc_state1 = self.getStateFromContext(self._ncmc_sim.context, self._state_keys)
+
         self._setStateTable('ncmc', 'state1', ncmc_state1)
 
     def _computeAlchemicalCorrection(self):
@@ -1114,6 +1119,18 @@ def _computeAlchemicalCorrection(self):
             -1.0 / self._ncmc_sim.context._integrator.kT)
 
         return correction_factor
+    def _checkPEConsistencyOnRejection(self):
+        """Checks wheter the potential energy is the same as the last MD iteration if a move is rejected
+        (it should be)
+        """
+        md_state0 = self.stateTable['md']['state0']
+        md_PE = self._md_sim.context.getState(getEnergy=True).getPotentialEnergy()
+        if not math.isclose(md_state0['potential_energy']._value, md_PE._value, rel_tol=float('1e-%s' % rtol)):
+            logger.error(
+                'Last MD potential energy %s != Current MD potential energy %s. Potential energy should match the prior state.'
+                % (md_state0['potential_energy'], md_PE))
+            sys.exit(1)
+
 
     def _acceptRejectMove(self, write_move=False):
         """Choose to accept or reject the proposed move based
@@ -1126,6 +1143,7 @@ def _acceptRejectMove(self, write_move=False):
         """
         work_ncmc = self._ncmc_sim.context._integrator.getLogAcceptanceProbability(self._ncmc_sim.context)
         randnum = math.log(np.random.random())
+        acceptance_ratio = self._move_engine.selected_move.acceptance_ratio
 
         # Compute correction if work_ncmc is not NaN
         if not np.isnan(work_ncmc):
@@ -1134,7 +1152,14 @@ def _acceptRejectMove(self, write_move=False):
                 'NCMCLogAcceptanceProbability = %.6f + Alchemical Correction = %.6f' % (work_ncmc, correction_factor))
             work_ncmc = work_ncmc + correction_factor
 
-        if work_ncmc > randnum:
+        if np.isclose(0, acceptance_ratio):
+            print('accept probablility', acceptance_ratio)
+            logger.info('NCMC MOVE REJECTED: acceptance_ratio for the selected move is 0; work_ncmc {}'.format(work_ncmc))
+            self._checkPEConsistencyOnRejection()
+
+
+
+        elif work_ncmc + math.log(acceptance_ratio) > randnum:
             self.accept += 1
             logger.info('NCMC MOVE ACCEPTED: work_ncmc {} > randnum {}'.format(work_ncmc, randnum))
 
@@ -1152,15 +1177,8 @@ def _acceptRejectMove(self, write_move=False):
 
             # If reject move, do nothing,
             # NCMC simulation be updated from MD Simulation next iteration.
-
             # Potential energy should be from last MD step in the previous iteration
-            md_state0 = self.stateTable['md']['state0']
-            md_PE = self._md_sim.context.getState(getEnergy=True).getPotentialEnergy()
-            if not math.isclose(md_state0['potential_energy']._value, md_PE._value, rel_tol=float('1e-%s' % rtol)):
-                logger.error(
-                    'Last MD potential energy %s != Current MD potential energy %s. Potential energy should match the prior state.'
-                    % (md_state0['potential_energy'], md_PE))
-                sys.exit(1)
+            self._checkPEConsistencyOnRejection()
 
     def _resetSimulations(self, temperature=None):
         """At the end of each iteration:
@@ -1287,8 +1305,14 @@ def _acceptRejectMove(self, temperature=None):
         work_mc = (md_state1['potential_energy'] - md_state0['potential_energy']) * (
             -1.0 / self._ncmc_sim.context._integrator.kT)
         randnum = math.log(np.random.random())
+        acceptance_ratio = self._move_engine.selected_move.acceptance_ratio
+
+        if np.isclose(0, acceptance_ratio):
+            self.reject += 1
+            logger.info('MC MOVE REJECTED: work_mc {} < {}'.format(work_mc, randnum))
+            self._md_sim.context.setPositions(md_state0['positions'])
 
-        if work_mc > randnum:
+        elif work_mc > randnum:
             self.accept += 1
             logger.info('MC MOVE ACCEPTED: work_mc {} > randnum {}'.format(work_mc, randnum))
             self._md_sim.context.setPositions(md_state1['positions'])

blues/tests/rotpref.yml

@@ -0,0 +1,58 @@
+ARG:
+    1: [-68, 60, -171]
+    2: [-67, 70, 182]
+    3: [-66, 68, -178]
+    4: [-88, 86, -179]
+    5: [-60, 60, 180] 
+ASN:
+    1: [-68, 59, 180]
+    2: [-107, 53] 
+ASP:
+    1: [-75, 74, -182]
+    2: [107, 62, -171]
+CYS:
+    1: [-62, 55, -176]
+GLN:
+    1: [-67, 56, -171]
+    2: [-66, 82, -175]
+    3: [-85, 90]
+GLU:
+    1: [-179,-68,60]
+    2: [-179,-70,85]
+    3: [-68,88]
+HIS:
+    1: [-60, 58, -179]
+    2: [-80, 100]
+ILE:
+    1: [-60, 50, 170]
+    2: [-60, 70, -178]
+LEU:
+    1: [-60, 180]
+    2: [68, 180]
+LYS:
+    1: [-68, 60, -178]
+    2: [-68, 56, -178]
+    3: [-68, 56, -178]
+    4: [-68, 60, -178]
+    5: [-68, 56, -178]
+MET: 
+    1: [-70, 60, 180]
+    2: [-60, 60, 180]
+    3: [-70, 70, 175]
+PHE:
+    1: [-79, 52, -177]
+    2: [-90, 90]
+PRO:
+    1: [-30, 30]
+SER:
+    1: [-60, 60, -171]
+THR:
+    1: [-60, 55, 170]
+TRP:
+    1: [-68, 68, 178]
+    2: [-90, 0, 90]
+TYR:
+    1: [-68, 68, 178]
+    2: [-90, 90]
+VAL:
+    1: [-65,65,178]

blues/tests/test_sidechain.py

@@ -24,7 +24,88 @@
 
 
 @skipUnless(has_openeye, 'Cannot test SideChainMove without openeye-toolkits and oeommtools.')
-class SideChainTester(unittest.TestCase):
+class SideChainTesterDunbrak(unittest.TestCase):
+    """
+    Test the SmartDartMove.move() function.
+    """
+
+    def setUp(self):
+        # Obtain topologies/positions
+        prmtop = utils.get_data_filename('blues', 'tests/data/vacDivaline.prmtop')
+        inpcrd = utils.get_data_filename('blues', 'tests/data/vacDivaline.inpcrd')
+        self.struct = parmed.load_file(prmtop, xyz=inpcrd)
+
+        self.sidechain = SideChainMove(self.struct, [1], bias_range=15.0)
+        self.engine = MoveEngine(self.sidechain)
+        self.engine.selectMove()
+
+        self.system_cfg = {'nonbondedMethod': app.NoCutoff, 'constraints': app.HBonds}
+        self.systems = SystemFactory(self.struct, self.sidechain.atom_indices, self.system_cfg)
+
+        self.cfg = {
+            'dt': 0.002 * unit.picoseconds,
+            'friction': 1 * 1 / unit.picoseconds,
+            'temperature': 300 * unit.kelvin,
+            'nIter': 1,
+            'nstepsMD': 1,
+            'nstepsNC': 4,
+            'alchemical_functions': {
+                'lambda_sterics':
+                'step(0.199999-lambda) + step(lambda-0.2)*step(0.8-lambda)*abs(lambda-0.5)*1/0.3 + step(lambda-0.800001)',
+                'lambda_electrostatics':
+                'step(0.2-lambda)- 1/0.2*lambda*step(0.2-lambda) + 1/0.2*(lambda-0.8)*step(lambda-0.8)'
+            }
+        }
+
+        self.simulations = SimulationFactory(self.systems, self.engine, self.cfg)
+
+    def test_getRotBondAtoms(self):
+        vals = [v for v in self.sidechain.rot_atoms[1]['chis'][1]['atms2mv']]
+        assert len(vals) == 11
+        #Ensure it selects 1 rotatable bond in Valine
+        assert len(self.sidechain.rot_bonds) == 1
+
+    def test_sidechain_move(self):
+        #check to make sure the dart stays within the specified ranges during the simulation
+        atom_indices = [v for v in self.sidechain.rot_atoms[1]['chis'][1]['atms2mv']]
+
+        for i in range(20):
+            before_context = self.simulations.ncmc.context
+            dobefore = self.sidechain.beforeMove(before_context)
+            before_move = self.simulations.ncmc.context.getState(getPositions=True).getPositions(
+                asNumpy=True)[atom_indices, :]
+            self.simulations.ncmc.context = self.engine.runEngine(self.simulations.ncmc.context)
+            after_move = self.simulations.ncmc.context.getState(getPositions=True).getPositions(
+                asNumpy=True)[atom_indices, :]
+
+            #Check that our system has performed the move correctly
+            # Integrator must step for context to update positions
+            # Remove the first two atoms in check as these are the anchor atoms and are not rotated.
+            pos_compare = np.not_equal(before_move, after_move)[2:, :].all()
+            positions = self.simulations.ncmc.context.getState(getPositions=True).getPositions(asNumpy=True)
+            bond_check = False
+            bonds_in_bins = []
+            #check the current angle after the move and make sure it's within range of the dihedral bins specified
+            for residx in self.sidechain.rot_atoms.keys():
+                resname = self.sidechain.rot_atoms[residx]['res_name']
+                for chi in self.sidechain.rot_atoms[residx]['chis']:
+                    dihed_atoms = [self.sidechain.rot_atoms[residx]['chis'][chi]['dihed_atms']]
+                    curr_angle = self.sidechain.getDihedral(positions,dihed_atoms)
+                    bin_ct=0
+                    for bin in self.sidechain.rot_atoms[residx]['chis'][chi]['bin_pref']:
+                        bin_ct+=1
+                        if self.sidechain.is_in_bin(curr_angle,bin):
+
+                            bin_idx = bin_ct-1
+                            bonds_in_bins.append([(self.sidechain.rot_atoms[residx]['chis'][chi]),curr_angle[0][0],bin_idx])
+                            bond_check = True
+                            break
+
+            assert bond_check == True
+
+            assert pos_compare
+
+class SideChainTesterRandom(unittest.TestCase):
     """
     Test the SmartDartMove.move() function.
     """
@@ -60,13 +141,15 @@ def setUp(self):
         self.simulations = SimulationFactory(self.systems, self.engine, self.cfg)
 
     def test_getRotBondAtoms(self):
-        vals = [v for v in self.sidechain.rot_atoms[1].values()][0]
+        vals = [v for v in self.sidechain.rot_atoms[1]['chis'][1]['atms2mv']]
         assert len(vals) == 11
         #Ensure it selects 1 rotatable bond in Valine
         assert len(self.sidechain.rot_bonds) == 1
 
     def test_sidechain_move(self):
-        atom_indices = [v for v in self.sidechain.rot_atoms[1].values()][0]
+        atom_indices = [v for v in self.sidechain.rot_atoms[1]['chis'][1]['atms2mv']]
+        before_context = self.simulations.ncmc.context
+        dobefore = self.sidechain.beforeMove(before_context)
         before_move = self.simulations.ncmc.context.getState(getPositions=True).getPositions(
             asNumpy=True)[atom_indices, :]
         self.simulations.ncmc.context = self.engine.runEngine(self.simulations.ncmc.context)
@@ -79,6 +162,5 @@ def test_sidechain_move(self):
         pos_compare = np.not_equal(before_move, after_move)[2:, :].all()
         assert pos_compare
 
-
 if __name__ == "__main__":
     unittest.main()

blues/tests/test_simulation.py

@@ -78,6 +78,21 @@ def system(structure, system_cfg):
 
 
 class NoRandomLigandRotation(RandomLigandRotationMove):
+    def __init__(self, structure, resname='LIG'):
+        print('type structure', type(structure), structure)
+        super(NoRandomLigandRotation, self).__init__(structure, resname)
+        self.before = False
+        self.after= False
+        self.acceptance_ratio = 1
+    def beforeMove(self, context):
+        if self.before:
+            self.acceptance_ratio = 0
+        return context
+    def afterMove(self, context):
+        if self.after:
+            self.acceptance_ratio = 0
+        return context
+
     def move(self, context):
         return context
 
@@ -493,7 +508,13 @@ def test_blues_simulationRunYAML(self, tmpdir, structure, tol_atom_indices, syst
         pos_compare = np.not_equal(before_iter, after_iter).all()
         assert pos_compare
 
-    def test_blues_simulationRunPython(self, systems, simulations, engine, tmpdir, sim_cfg):
+    @pytest.mark.parametrize('before,after', [
+        (True, False),
+        (False, True),
+        (True, True),
+        (False, False)
+        ])
+    def test_blues_simulationRunPython(self, systems, simulations, engine, tmpdir, sim_cfg, before, after):
         print('Testing BLUESSimulation.run() from pure python')
         md_rep_cfg = {
             'stream': {

examples/example_sidechain.py

@@ -1,11 +1,11 @@
 from blues.moves import SideChainMove
-from blues.engine import MoveEngine
+from blues.moves import MoveEngine
 from blues.simulation import *
 import json
 from blues.settings import *
 
 # Parse a YAML configuration, return as Dict
-cfg = Settings('sidechain_cuda.yaml').asDict()
+cfg = Settings('sidechain_cuda.yml').asDict()
 structure = cfg['Structure']
 
 #Select move type
@@ -28,7 +28,7 @@
 import mdtraj as md
 import numpy as np
 
-traj = md.load_netcdf('vacDivaline-test/vacDivaline.nc', top='tests/data/vacDivaline.prmtop')
+traj = md.load_netcdf('vacDivaline-test/vacDivaline.nc', top='../blues/tests/data/vacDivaline.prmtop')
 indicies = np.array([[0, 4, 6, 8]])
 dihedraldata = md.compute_dihedrals(traj, indicies)
 with open("vacDivaline-test/dihedrals.txt", 'w') as output: