diff --git a/perses/rjmc/geometry.py b/perses/rjmc/geometry.py index e897a9115..fd46868bc 100644 --- a/perses/rjmc/geometry.py +++ b/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 diff --git a/perses/tests/test_geometry_engine.py b/perses/tests/test_geometry_engine.py index 52097bb28..eefa96b41 100644 --- a/perses/tests/test_geometry_engine.py +++ b/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,7 +298,7 @@ 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) @@ -306,7 +306,7 @@ def test_logp_reverse(): #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()