Adding constraints, fixing a couple tests.

perses/rjmc/geometry.py

@@ -8,6 +8,8 @@
 import numpy as np
 import copy
 from perses.rjmc import coordinate_tools
+import simtk.openmm as openmm
+
 
 
 
@@ -119,12 +121,18 @@ def propose(self, top_proposal, current_positions, beta):
                     torsion_atom = torsion.atom1
 
                 #propose a bond and calculate its probability
+                #if it's not a bond, it's a constraint
                 bond = self._get_relevant_bond(atom, bond_atom)
-                r_proposed = self._propose_bond(bond, beta)
-                bond_k = bond.type.k
-                sigma_r = units.sqrt(1/(beta*bond_k))
-                logZ_r = np.log((np.sqrt(2*np.pi)*sigma_r/sigma_r.unit))
-                logp_r = self._bond_logq(r_proposed, bond, beta) - logZ_r
+                if bond is not None:
+                    r_proposed = self._propose_bond(bond, beta)
+                    bond_k = bond.type.k
+                    sigma_r = units.sqrt(1/(beta*bond_k))
+                    logZ_r = np.log((np.sqrt(2*np.pi)*sigma_r/sigma_r.unit))
+                    logp_r = self._bond_logq(r_proposed, bond, beta) - logZ_r
+                else:
+                    constraint = self._get_bond_constraint(atom, bond_atom, top_proposal.new_system)
+                    r_proposed = constraint #set bond length to exactly constraint
+                    logp_r = 0.0
 
                 #propose an angle and calculate its probability
                 angle = self._get_relevant_angle(atom, bond_atom, angle_atom)
@@ -208,13 +216,19 @@ def logp_reverse(self, top_proposal, new_coordinates, old_coordinates, beta):
                 theta = internal_coordinates[1]*units.radian
                 phi = internal_coordinates[2]*units.radian
 
-                #propose a bond and calculate its probability
+                #propose a bond and calculate its probability, if it's a bond.
+                #Otherwise, it's a constraint.
                 bond = self._get_relevant_bond(atom, bond_atom)
-                bond_k = bond.type.k
-                sigma_r = units.sqrt(1/(beta*bond_k))
-                logZ_r = np.log((np.sqrt(2*np.pi)*sigma_r/sigma_r.unit)) #need to eliminate unit to allow numpy log
-                logp_r = self._bond_logq(r, bond, beta) - logZ_r
-
+                if bond is not None:
+                    bond_k = bond.type.k
+                    sigma_r = units.sqrt(1/(beta*bond_k))
+                    logZ_r = np.log((np.sqrt(2*np.pi)*sigma_r/sigma_r.unit)) #need to eliminate unit to allow numpy log
+                    logp_r = self._bond_logq(r, bond, beta) - logZ_r
+                else:
+                    constraint = self._get_bond_constraint(atom, bond_atom, top_proposal.old_system)
+                    if constraint is None:
+                        raise Exception("No Bond or constraint between atom %d and %d" % (atom.idx, bond_atom.idx))
+                    logp_r = 0.0
                 #propose an angle and calculate its probability
                 angle = self._get_relevant_angle(atom, bond_atom, angle_atom)
                 angle_k = angle.type.k
@@ -230,26 +244,11 @@ def logp_reverse(self, top_proposal, new_coordinates, old_coordinates, beta):
                 old_unique_atoms.remove(atom)
         return logp
 
-    def _get_proposal_order(self, topology_for_proposal, reference_topology, new_to_old_atom_map):
-        """
-        Get the order of proposal for the atoms needing proposal. Also get a list of torsions
-        that could be used in each atom's proposal.
-
-        Parameters
-        ----------
-        topology_for_proposal
-        reference_topology
-        new_to_old_atom_map
-
-        Returns
-        -------
-        atoms_for_proposal : dict of {atom :
-
-        """
-
     def _get_relevant_bond(self, atom1, atom2):
         """
-        utility function to get the bond connecting atoms 1 and 2
+        utility function to get the bond connecting atoms 1 and 2.
+        Returns either a bond object or None
+        (since there is no constraint class)
 
         Arguments
         ---------
@@ -261,15 +260,50 @@ def _get_relevant_bond(self, atom1, atom2):
         Returns
         -------
         bond : bond object
-            Bond connecting the two atoms
+            Bond connecting the two atoms, if there is one. None if constrained or
+            no bond.
         """
         bonds_1 = set(atom1.bonds)
         bonds_2 = set(atom2.bonds)
         relevant_bond_set = bonds_1.intersection(bonds_2)
         relevant_bond = relevant_bond_set.pop()
+        if relevant_bond.type is None:
+            return None
         relevant_bond_with_units = self._add_bond_units(relevant_bond)
         return relevant_bond_with_units
 
+    def _get_bond_constraint(self, atom1, atom2, system):
+        """
+        Get the constraint parameters corresponding to the bond
+        between the given atoms
+
+        Parameters
+        ----------
+        atom1 : parmed.Atom object
+           the first atom of the constrained bond
+        atom2 : parmed.Atom object
+           the second atom of the constrained bond
+        system : openmm.System object
+           The system containing the constraint
+
+        Returns
+        -------
+        constraint : float, quantity nm
+            the parameters of the bond constraint
+        """
+        atom_indices = {atom1.idx, atom2.idx}
+        n_constraints = system.getNumConstraints()
+        constraint = None
+        for i in range(n_constraints):
+            constraint_parameters = system.getConstraintParameters(i)
+            constraint_atoms = set(constraint_parameters[:2])
+            if len(constraint_atoms.intersection(atom_indices))==2:
+                constraint = constraint_parameters[2]
+        return constraint
+
+
+
+
     def _get_relevant_angle(self, atom1, atom2, atom3):
         """
         Get the angle containing the 3 given atoms

perses/tests/test_geometry_engine.py

@@ -67,7 +67,7 @@ def oemol_to_openmm_system(oemol, molecule_name):
     from perses.rjmc import topology_proposal
     from openmoltools import forcefield_generators
     gaff_xml_filename = get_data_filename('data/gaff.xml')
-    system_generator = topology_proposal.SystemGenerator([gaff_xml_filename])
+    system_generator = topology_proposal.SystemGenerator([gaff_xml_filename], forcefield_kwargs={'constraints' : app.HBonds})
     topology = forcefield_generators.generateTopologyFromOEMol(oemol)
     system = system_generator.build_system(topology)
     positions = extractPositionsFromOEMOL(oemol)
@@ -133,8 +133,8 @@ def test_run_geometry_engine():
     molecule2 = generate_initial_molecule(molecule_name_2)
     new_to_old_atom_mapping = align_molecules(molecule1, molecule2)
 
-    sys1, pos1, top1 = oemol_to_openmm_system_amber(molecule1, molecule_name_1)
-    sys2, pos2, top2 = oemol_to_openmm_system_amber(molecule2, molecule_name_2)
+    sys1, pos1, top1 = oemol_to_openmm_system(molecule1, molecule_name_1)
+    sys2, pos2, top2 = oemol_to_openmm_system(molecule2, molecule_name_2)
 
     import perses.rjmc.geometry as geometry
     import perses.rjmc.topology_proposal as topology_proposal
@@ -152,7 +152,7 @@ def test_run_geometry_engine():
     state = context.getState(getEnergy=True)
     print("Energy before proposal is: %s" % str(state.getPotentialEnergy()))
 
-    new_positions, logp_proposal = geometry_engine.propose(sm_top_proposal)
+    new_positions, logp_proposal = geometry_engine.propose(sm_top_proposal, pos1, beta)
     app.PDBFile.writeFile(top2, new_positions, file=test_pdb_file)
     test_pdb_file.close()
     context.setPositions(new_positions)
@@ -298,15 +298,15 @@ def test_logp_reverse():
     import perses.rjmc.topology_proposal as topology_proposal
 
     sm_top_proposal = topology_proposal.TopologyProposal(new_topology=top2, new_system=sys2, old_topology=top1, old_system=sys1,
-                                                                    logp_proposal=0.0, new_to_old_atom_map=new_to_old_atom_mapping, metadata={'test':0.0})
+                                                                    logp_proposal=0.0, new_to_old_atom_map=new_to_old_atom_mapping, new_chemical_state_key="CCC", old_chemical_state_key="CC", metadata={'test':0.0})
     geometry_engine = geometry.FFAllAngleGeometryEngine({'test': 'true'})
     new_positions, logp_proposal = geometry_engine.propose(sm_top_proposal, pos1, beta)
 
     #now pretend that the new coordinates are the old, and calculate logp again
     #reverse the atom map:
     old_to_new_atom_mapping = {value : key for key, value in new_to_old_atom_mapping.items()}
     sm_reverse_proposal = topology_proposal.TopologyProposal(new_topology=top1, new_system=sys1, old_topology=top2, old_system=sys2,
-                                                                      logp_proposal=0.0, new_to_old_atom_map=old_to_new_atom_mapping, metadata={'test':0.0})
+                                                                      logp_proposal=0.0, new_to_old_atom_map=old_to_new_atom_mapping, new_chemical_state_key="CC", old_chemical_state_key="CCCC", metadata={'test':0.0})
     logp_reverse = geometry_engine.logp_reverse(sm_reverse_proposal, pos1, new_positions, beta)
     print(logp_proposal)
     print(logp_reverse)
@@ -375,7 +375,7 @@ def test_angle():
 
 if __name__=="__main__":
     #test_coordinate_conversion()
-    #test_run_geometry_engine()
+    test_run_geometry_engine()
     #test_existing_coordinates()
     #test_openmm_dihedral()
     #test_try_random_itoc()