"""
This contains the base class for the geometry engine, which proposes new positions
for each additional atom that must be added.
"""
import parmed
import simtk.unit as units
import logging
import numpy as np
import copy
#from perses.rjmc import coordinate_numba
import simtk.openmm as openmm
import collections
import openeye.oechem as oechem
import openeye.oeomega as oeomega
import simtk.openmm.app as app
import time
import logging
from perses.storage import NetCDFStorage, NetCDFStorageView
_logger = logging.getLogger("geometry")
[docs]class GeometryEngine(object):
"""
This is the base class for the geometry engine.
Arguments
---------
metadata : dict
GeometryEngine-related metadata as a dict
"""
[docs] def __init__(self, metadata=None, storage=None):
# TODO: Either this base constructor should be called by subclasses, or we should remove its arguments.
pass
[docs] def propose(self, top_proposal, current_positions, beta):
"""
Make a geometry proposal for the appropriate atoms.
Arguments
----------
top_proposal : TopologyProposal object
Object containing the relevant results of a topology proposal
beta : float
The inverse temperature
Returns
-------
new_positions : [n, 3] ndarray
The new positions of the system
"""
return np.array([0.0,0.0,0.0])
[docs] def logp_reverse(self, top_proposal, new_coordinates, old_coordinates, beta):
"""
Calculate the logp for the given geometry proposal
Arguments
----------
top_proposal : TopologyProposal object
Object containing the relevant results of a topology proposal
new_coordinates : [n, 3] np.ndarray
The coordinates of the system after the proposal
old_coordiantes : [n, 3] np.ndarray
The coordinates of the system before the proposal
direction : str, either 'forward' or 'reverse'
whether the transformation is for the forward NCMC move or the reverse
beta : float
The inverse temperature
Returns
-------
logp : float
The log probability of the proposal for the given transformation
"""
return 0.0
[docs]class FFAllAngleGeometryEngine(GeometryEngine):
"""
This is an implementation of GeometryEngine which uses all valence terms and OpenMM
Parameters
----------
use_sterics : bool, optional, default=False
If True, sterics will be used in proposals to minimize clashes.
This may significantly slow down the simulation, however.
"""
[docs] def __init__(self, metadata=None, use_sterics=False, n_torsion_divisions=180, verbose=True, storage=None, bond_softening_constant=1.0, angle_softening_constant=1.0):
self._metadata = metadata
self.write_proposal_pdb = False # if True, will write PDB for sequential atom placements
self.pdb_filename_prefix = 'geometry-proposal' # PDB file prefix for writing sequential atom placements
self.nproposed = 0 # number of times self.propose() has been called
self._energy_time = 0.0
self._torsion_coordinate_time = 0.0
self._position_set_time = 0.0
self.verbose = verbose
self.use_sterics = use_sterics
self._n_torsion_divisions = n_torsion_divisions
self._bond_softening_constant = bond_softening_constant
self._angle_softening_constant = angle_softening_constant
if storage:
self._storage = NetCDFStorageView(modname="GeometryEngine", storage=storage)
else:
self._storage = None
[docs] def propose(self, top_proposal, current_positions, beta):
"""
Make a geometry proposal for the appropriate atoms.
Arguments
----------
top_proposal : TopologyProposal object
Object containing the relevant results of a topology proposal
beta : float
The inverse temperature
Returns
-------
new_positions : [n, 3] ndarray
The new positions of the system
logp_proposal : float
The log probability of the forward-only proposal
"""
current_positions = current_positions.in_units_of(units.nanometers)
if not top_proposal.unique_new_atoms:
structure = parmed.openmm.load_topology(top_proposal.old_topology, top_proposal.old_system)
atoms_with_positions = [structure.atoms[atom_idx] for atom_idx in top_proposal.new_to_old_atom_map.keys()]
new_positions = self._copy_positions(atoms_with_positions, top_proposal, current_positions)
return new_positions, 0.0
logp_proposal, new_positions = self._logp_propose(top_proposal, current_positions, beta, direction='forward')
self.nproposed += 1
return new_positions, logp_proposal
[docs] def logp_reverse(self, top_proposal, new_coordinates, old_coordinates, beta):
"""
Calculate the logp for the given geometry proposal
Arguments
----------
top_proposal : TopologyProposal object
Object containing the relevant results of a topology proposal
new_coordinates : [n, 3] np.ndarray
The coordinates of the system after the proposal
old_coordiantes : [n, 3] np.ndarray
The coordinates of the system before the proposal
beta : float
The inverse temperature
Returns
-------
logp : float
The log probability of the proposal for the given transformation
"""
if not top_proposal.unique_old_atoms:
return 0.0
new_coordinates = new_coordinates.in_units_of(units.nanometers)
old_coordinates = old_coordinates.in_units_of(units.nanometers)
logp_proposal, _ = self._logp_propose(top_proposal, old_coordinates, beta, new_positions=new_coordinates, direction='reverse')
return logp_proposal
def _write_partial_pdb(self, pdbfile, topology, positions, atoms_with_positions, model_index):
"""
Write the subset of the molecule for which positions are defined.
"""
from simtk.openmm.app import Modeller
modeller = Modeller(topology, positions)
atom_indices_with_positions = [ atom.idx for atom in atoms_with_positions ]
atoms_to_delete = [ atom for atom in modeller.topology.atoms() if (atom.index not in atom_indices_with_positions) ]
modeller.delete(atoms_to_delete)
pdbfile.write('MODEL %5d\n' % model_index)
from simtk.openmm.app import PDBFile
PDBFile.writeFile(modeller.topology, modeller.positions, file=pdbfile)
pdbfile.flush()
pdbfile.write('ENDMDL\n')
def _logp_propose(self, top_proposal, old_positions, beta, new_positions=None, direction='forward'):
"""
This is an INTERNAL function that handles both the proposal and the logp calculation,
to reduce code duplication. Whether it proposes or just calculates a logp is based on
the direction option. Note that with respect to "new" and "old" terms, "new" will always
mean the direction we are proposing (even in the reverse case), so that for a reverse proposal,
this function will still take the new coordinates as new_coordinates
Parameters
----------
top_proposal : topology_proposal.TopologyProposal object
topology proposal containing the relevant information
old_positions : np.ndarray [n,3] in nm
The old coordinates.
beta : float
Inverse temperature
new_positions : np.ndarray [n,3] in nm, optional for forward
The new coordinates, if any. For proposal this is none
direction : str
Whether to make a proposal (forward) or just calculate logp (reverse)
Returns
-------
logp_proposal : float
the logp of the proposal
new_positions : [n,3] np.ndarray
The new positions (same as input if direction='reverse')
"""
initial_time = time.time()
proposal_order_tool = ProposalOrderTools(top_proposal)
proposal_order_time = time.time() - initial_time
growth_parameter_name = 'growth_stage'
if direction=="forward":
forward_init = time.time()
atom_proposal_order, logp_choice = proposal_order_tool.determine_proposal_order(direction='forward')
proposal_order_forward = time.time() - forward_init
structure = parmed.openmm.load_topology(top_proposal.new_topology, top_proposal.new_system)
#find and copy known positions
atoms_with_positions = [structure.atoms[atom_idx] for atom_idx in top_proposal.new_to_old_atom_map.keys()]
new_positions = self._copy_positions(atoms_with_positions, top_proposal, old_positions)
system_init = time.time()
growth_system_generator = GeometrySystemGenerator(top_proposal.new_system, atom_proposal_order.keys(), growth_parameter_name, reference_topology=top_proposal.new_topology, use_sterics=self.use_sterics)
growth_system = growth_system_generator.get_modified_system()
growth_system_time = time.time() - system_init
elif direction=='reverse':
if new_positions is None:
raise ValueError("For reverse proposals, new_positions must not be none.")
atom_proposal_order, logp_choice = proposal_order_tool.determine_proposal_order(direction='reverse')
structure = parmed.openmm.load_topology(top_proposal.old_topology, top_proposal.old_system)
atoms_with_positions = [structure.atoms[atom_idx] for atom_idx in top_proposal.old_to_new_atom_map.keys()]
growth_system_generator = GeometrySystemGenerator(top_proposal.old_system, atom_proposal_order.keys(), growth_parameter_name, reference_topology=top_proposal.old_topology, use_sterics=self.use_sterics)
growth_system = growth_system_generator.get_modified_system()
else:
raise ValueError("Parameter 'direction' must be forward or reverse")
logp_proposal = logp_choice
if self._storage:
self._storage.write_object("{}_proposal_order".format(direction), proposal_order_tool, iteration=self.nproposed)
if self.use_sterics:
platform_name = 'CPU'
else:
platform_name = 'Reference'
platform = openmm.Platform.getPlatformByName(platform_name)
integrator = openmm.VerletIntegrator(1*units.femtoseconds)
context = openmm.Context(growth_system, integrator, platform)
growth_system_generator.set_growth_parameter_index(len(atom_proposal_order.keys())+1, context)
growth_parameter_value = 1
#now for the main loop:
logging.debug("There are %d new atoms" % len(atom_proposal_order.items()))
atom_placements = []
for atom, torsion in atom_proposal_order.items():
growth_system_generator.set_growth_parameter_index(growth_parameter_value, context=context)
bond_atom = torsion.atom2
angle_atom = torsion.atom3
torsion_atom = torsion.atom4
if self.verbose: _logger.info("Proposing atom %s from torsion %s" %(str(atom), str(torsion)))
if atom != torsion.atom1:
raise Exception('atom != torsion.atom1')
#get internal coordinates if direction is reverse
if direction=='reverse':
atom_coords = old_positions[atom.idx]
bond_coords = old_positions[bond_atom.idx]
angle_coords = old_positions[angle_atom.idx]
torsion_coords = old_positions[torsion_atom.idx]
internal_coordinates, detJ = self._cartesian_to_internal(atom_coords, bond_coords, angle_coords, torsion_coords)
r = internal_coordinates[0]*atom_coords.unit
theta = internal_coordinates[1]*units.radian
phi = internal_coordinates[2]*units.radian
bond = self._get_relevant_bond(atom, bond_atom)
if bond is not None:
if direction=='forward':
r = self._propose_bond(bond, beta)
bond_k = bond.type.k * self._bond_softening_constant
sigma_r = units.sqrt(1/(beta*bond_k))
logZ_r = np.log((np.sqrt(2*np.pi)*(sigma_r.value_in_unit(units.angstrom))))
logp_r = self._bond_logq(r, bond, beta) - logZ_r
else:
if direction == 'forward':
constraint = self._get_bond_constraint(atom, bond_atom, top_proposal.new_system)
if constraint is None:
raise ValueError("Structure contains a topological bond [%s - %s] with no constraint or bond information." % (str(atom), str(bond_atom)))
r = 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)
if direction=='forward':
theta = self._propose_angle(angle, beta)
angle_k = angle.type.k * self._angle_softening_constant
sigma_theta = units.sqrt(1/(beta*angle_k))
logZ_theta = np.log((np.sqrt(2*np.pi)*(sigma_theta.value_in_unit(units.radians))))
logp_theta = self._angle_logq(theta, angle, beta) - logZ_theta
#propose a torsion angle and calcualate its probability
if direction=='forward':
phi, logp_phi = self._propose_torsion(context, torsion, new_positions, r, theta, beta, n_divisions=self._n_torsion_divisions)
xyz, detJ = self._internal_to_cartesian(new_positions[bond_atom.idx], new_positions[angle_atom.idx], new_positions[torsion_atom.idx], r, theta, phi)
new_positions[atom.idx] = xyz
else:
old_positions_for_torsion = copy.deepcopy(old_positions)
logp_phi = self._torsion_logp(context, torsion, old_positions_for_torsion, r, theta, phi, beta, n_divisions=self._n_torsion_divisions)
#accumulate logp
#if direction == 'reverse':
if self.verbose: _logger.info('%8d logp_r %12.3f | logp_theta %12.3f | logp_phi %12.3f | log(detJ) %12.3f' % (atom.idx, logp_r, logp_theta, logp_phi, np.log(detJ)))
atom_placement_array = np.array([atom.idx, r.value_in_unit_system(units.md_unit_system),
theta.value_in_unit_system(units.md_unit_system),
phi.value_in_unit_system(units.md_unit_system),
logp_r, logp_theta, logp_phi, np.log(detJ)])
atom_placements.append(atom_placement_array)
logp_proposal += logp_r + logp_theta + logp_phi + np.log(detJ)
growth_parameter_value += 1
# DEBUG: Write PDB file for placed atoms
atoms_with_positions.append(atom)
total_time = time.time() - initial_time
#use a new array for each placement, since the variable size will be different.
if self._storage:
self._storage.write_array("atom_placement_logp_{}_{}".format(direction, self.nproposed), np.stack(atom_placements))
if direction=='forward':
logging.log(logging.DEBUG, "Proposal order time: %f s | Growth system generation: %f s | Total torsion scan time %f s | Total energy computation time %f s | Position set time %f s| Total time %f s" % (proposal_order_time, growth_system_time , self._torsion_coordinate_time, self._energy_time, self._position_set_time, total_time))
self._torsion_coordinate_time = 0.0
self._energy_time = 0.0
self._position_set_time = 0.0
return logp_proposal, new_positions
@staticmethod
def _oemol_from_residue(res, verbose=True):
"""
Get an OEMol from a residue, even if that residue
is polymeric. In the latter case, external bonds
are replaced by hydrogens.
Parameters
----------
res : app.Residue
The residue in question
verbose : bool, optional, default=False
If True, will print verbose output.
Returns
-------
oemol : openeye.oechem.OEMol
an oemol representation of the residue with topology indices
"""
# TODO: This seems to be broken. Can we fix it?
from openmoltools.forcefield_generators import generateOEMolFromTopologyResidue
external_bonds = list(res.external_bonds())
for bond in external_bonds:
if verbose: print(bond)
new_atoms = {}
highest_index = 0
if external_bonds:
new_topology = app.Topology()
new_chain = new_topology.addChain(0)
new_res = new_topology.addResidue("new_res", new_chain)
for atom in res.atoms():
new_atom = new_topology.addAtom(atom.name, atom.element, new_res, atom.id)
new_atom.index = atom.index
new_atoms[atom] = new_atom
highest_index = max(highest_index, atom.index)
for bond in res.internal_bonds():
new_topology.addBond(new_atoms[bond[0]], new_atoms[bond[1]])
for bond in res.external_bonds():
internal_atom = [atom for atom in bond if atom.residue==res][0]
if verbose:
print('internal atom')
print(internal_atom)
highest_index += 1
if internal_atom.name=='N':
if verbose: print('Adding H to N')
new_atom = new_topology.addAtom("H2", app.Element.getByAtomicNumber(1), new_res, -1)
new_atom.index = -1
new_topology.addBond(new_atoms[internal_atom], new_atom)
if internal_atom.name=='C':
if verbose: print('Adding OH to C')
new_atom = new_topology.addAtom("O2", app.Element.getByAtomicNumber(8), new_res, -1)
new_atom.index = -1
new_topology.addBond(new_atoms[internal_atom], new_atom)
highest_index += 1
new_hydrogen = new_topology.addAtom("HO", app.Element.getByAtomicNumber(1), new_res, -1)
new_hydrogen.index = -1
new_topology.addBond(new_hydrogen, new_atom)
res_to_use = new_res
external_bonds = list(res_to_use.external_bonds())
else:
res_to_use = res
oemol = generateOEMolFromTopologyResidue(res_to_use, geometry=False)
oechem.OEAddExplicitHydrogens(oemol)
return oemol
def _copy_positions(self, atoms_with_positions, top_proposal, current_positions):
"""
Copy the current positions to an array that will also hold new positions
Parameters
----------
atoms_with_positions : list of parmed.Atom
atoms that currently have positions
top_proposal : topology_proposal.TopologyProposal
topology proposal object
current_positions : [n, 3] np.ndarray in nm
Positions of the current system
Returns
-------
new_positions : np.ndarray in nm
Array for new positions with known positions filled in
"""
new_positions = units.Quantity(np.zeros([top_proposal.n_atoms_new, 3]), unit=units.nanometers)
# Workaround for CustomAngleForce NaNs: Create random non-zero positions for new atoms.
new_positions = units.Quantity(np.random.random([top_proposal.n_atoms_new, 3]), unit=units.nanometers)
current_positions = current_positions.in_units_of(units.nanometers)
#copy positions
for atom in atoms_with_positions:
old_index = top_proposal.new_to_old_atom_map[atom.idx]
new_positions[atom.idx] = current_positions[old_index]
return new_positions
def _get_relevant_bond(self, atom1, atom2):
"""
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
---------
atom1 : parmed atom object
One of the atoms in the bond
atom2 : parmed.atom object
The other atom in the bond
Returns
-------
bond : bond object
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
"""
atom1_angles = set(atom1.angles)
atom2_angles = set(atom2.angles)
atom3_angles = set(atom3.angles)
relevant_angle_set = atom1_angles.intersection(atom2_angles, atom3_angles)
# DEBUG
if len(relevant_angle_set) == 0:
print('atom1_angles:')
print(atom1_angles)
print('atom2_angles:')
print(atom2_angles)
print('atom3_angles:')
print(atom3_angles)
raise Exception('Atoms %s-%s-%s do not share a parmed Angle term' % (atom1, atom2, atom3))
relevant_angle = relevant_angle_set.pop()
if type(relevant_angle.type.k) != units.Quantity:
relevant_angle_with_units = self._add_angle_units(relevant_angle)
else:
relevant_angle_with_units = relevant_angle
return relevant_angle_with_units
def _add_bond_units(self, bond):
"""
Add the correct units to a harmonic bond
Arguments
---------
bond : parmed bond object
The bond to get units
Returns
-------
"""
if type(bond.type.k)==units.Quantity:
return bond
bond.type.req = units.Quantity(bond.type.req, unit=units.angstrom)
bond.type.k = units.Quantity(2.0*bond.type.k, unit=units.kilocalorie_per_mole/units.angstrom**2)
return bond
def _add_angle_units(self, angle):
"""
Add the correct units to a harmonic angle
Arguments
----------
angle : parmed angle object
the angle to get unit-ed
Returns
-------
angle_with_units : parmed angle
The angle, but with units on its parameters
"""
if type(angle.type.k)==units.Quantity:
return angle
angle.type.theteq = units.Quantity(angle.type.theteq, unit=units.degree)
angle.type.k = units.Quantity(2.0*angle.type.k, unit=units.kilocalorie_per_mole/units.radian**2)
return angle
def _add_torsion_units(self, torsion):
"""
Add the correct units to a torsion
Arguments
---------
torsion : parmed.dihedral object
The torsion needing units
Returns
-------
torsion : parmed.dihedral object
Torsion but with units added
"""
if type(torsion.type.phi_k) == units.Quantity:
return torsion
torsion.type.phi_k = units.Quantity(torsion.type.phi_k, unit=units.kilocalorie_per_mole)
torsion.type.phase = units.Quantity(torsion.type.phase, unit=units.degree)
return torsion
def _rotation_matrix(self, axis, angle):
"""
This method produces a rotation matrix given an axis and an angle.
"""
axis = axis/np.linalg.norm(axis)
axis_squared = np.square(axis)
cos_angle = np.cos(angle)
sin_angle = np.sin(angle)
rot_matrix_row_one = np.array([cos_angle+axis_squared[0]*(1-cos_angle),
axis[0]*axis[1]*(1-cos_angle) - axis[2]*sin_angle,
axis[0]*axis[2]*(1-cos_angle)+axis[1]*sin_angle])
rot_matrix_row_two = np.array([axis[1]*axis[0]*(1-cos_angle)+axis[2]*sin_angle,
cos_angle+axis_squared[1]*(1-cos_angle),
axis[1]*axis[2]*(1-cos_angle) - axis[0]*sin_angle])
rot_matrix_row_three = np.array([axis[2]*axis[0]*(1-cos_angle)-axis[1]*sin_angle,
axis[2]*axis[1]*(1-cos_angle)+axis[0]*sin_angle,
cos_angle+axis_squared[2]*(1-cos_angle)])
rotation_matrix = np.array([rot_matrix_row_one, rot_matrix_row_two, rot_matrix_row_three])
return rotation_matrix
def _cartesian_to_internal(self, atom_position, bond_position, angle_position, torsion_position):
"""
Cartesian to internal function
"""
from perses.rjmc import coordinate_numba
#ensure we have the correct units, then remove them
atom_position = atom_position.value_in_unit(units.nanometers).astype(np.float64)
bond_position = bond_position.value_in_unit(units.nanometers).astype(np.float64)
angle_position = angle_position.value_in_unit(units.nanometers).astype(np.float64)
torsion_position = torsion_position.value_in_unit(units.nanometers).astype(np.float64)
internal_coords = coordinate_numba.cartesian_to_internal(atom_position, bond_position, angle_position, torsion_position)
return internal_coords, np.abs(internal_coords[0]**2*np.sin(internal_coords[1]))
def _internal_to_cartesian(self, bond_position, angle_position, torsion_position, r, theta, phi):
"""
Calculate the cartesian coordinates given the internal, as well as abs(detJ)
"""
from perses.rjmc import coordinate_numba
r = r.value_in_unit(units.nanometers)
theta = theta.value_in_unit(units.radians)
phi = phi.value_in_unit(units.radians)
bond_position = bond_position.value_in_unit(units.nanometers).astype(np.float64)
angle_position = angle_position.value_in_unit(units.nanometers).astype(np.float64)
torsion_position = torsion_position.value_in_unit(units.nanometers).astype(np.float64)
xyz = coordinate_numba.internal_to_cartesian(bond_position, angle_position, torsion_position, np.array([r, theta, phi], dtype=np.float64))
xyz = units.Quantity(xyz, unit=units.nanometers)
return xyz, np.abs(r**2*np.sin(theta))
def _bond_logq(self, r, bond, beta):
"""
Calculate the log-probability of a given bond at a given inverse temperature
Arguments
---------
r : float
bond length, in nanometers
r0 : float
equilibrium bond length, in nanometers
k_eq : float
Spring constant of bond
beta : float
1/kT or inverse temperature
"""
k_eq = bond.type.k
r0 = bond.type.req
logq = -beta*0.5*k_eq*(r-r0)**2
return logq
def _angle_logq(self, theta, angle, beta):
"""
Calculate the log-probability of a given bond at a given inverse temperature
Arguments
---------
theta : float
bond angle, in randians
angle : parmed angle object
Bond angle object containing parameters
beta : float
1/kT or inverse temperature
"""
k_eq = angle.type.k
theta0 = angle.type.theteq
logq = -beta*k_eq*0.5*(theta-theta0)**2
return logq
def _propose_bond(self, bond, beta):
"""
Bond length proposal
"""
r0 = bond.type.req
k = bond.type.k * self._bond_softening_constant
sigma_r = units.sqrt(1.0/(beta*k))
r = sigma_r*np.random.randn() + r0
return r
def _propose_angle(self, angle, beta):
"""
Bond angle proposal
"""
theta0 = angle.type.theteq
k = angle.type.k * self._angle_softening_constant
sigma_theta = units.sqrt(1.0/(beta*k))
theta = sigma_theta*np.random.randn() + theta0
return theta
def _torsion_scan(self, torsion, positions, r, theta, n_divisions=360):
"""
Rotate the atom about the
Parameters
----------
torsion : parmed.Dihedral
parmed Dihedral containing relevant atoms
positions : [n,3] np.ndarray in nm
positions of the atoms in the system
r : float in nm
bond length
theta : float in radians
bond angle
Returns
-------
xyzs : np.ndarray, in nm
The cartesian coordinates of each
phis : np.ndarray, in radians
The torsions angles at which a potential will be calculated
"""
from perses.rjmc import coordinate_numba
torsion_scan_init = time.time()
positions_copy = copy.deepcopy(positions)
positions_copy = positions_copy.value_in_unit(units.nanometers)
positions_copy = positions_copy.astype(np.float64)
r = r.value_in_unit(units.nanometers)
theta = theta.value_in_unit(units.radians)
bond_atom = torsion.atom2
angle_atom = torsion.atom3
torsion_atom = torsion.atom4
phis = np.arange(-np.pi, +np.pi, (2.0*np.pi)/n_divisions) # Can't use units here.
xyzs = coordinate_numba.torsion_scan(positions_copy[bond_atom.idx], positions_copy[angle_atom.idx], positions_copy[torsion_atom.idx], np.array([r, theta, 0.0]), phis)
xyzs_quantity = units.Quantity(xyzs, unit=units.nanometers) #have to put the units back now
phis = units.Quantity(phis, unit=units.radians)
torsion_scan_time = time.time() - torsion_scan_init
self._torsion_coordinate_time += torsion_scan_time
return xyzs_quantity, phis
def _torsion_log_probability_mass_function(self, growth_context, torsion, positions, r, theta, beta, n_divisions=360):
"""
Calculate the torsion logp pmf using OpenMM
Parameters
----------
growth_context : openmm.Context
Context containing the modified system and
torsion : parmed.Dihedral
parmed Dihedral containing relevant atoms
positions : [n,3] np.ndarray in nm
positions of the atoms in the system
r : float in nm
bond length
theta : float in radians
bond angle
beta : float
inverse temperature
n_divisions : int, optional
number of divisions for the torsion scan
Returns
-------
logp_torsions : np.ndarray of float
normalized probability of each of n_divisions of torsion
phis : np.ndarray, in radians
The torsions angles at which a potential was calculated
"""
logq = np.zeros(n_divisions)
atom_idx = torsion.atom1.idx
xyzs, phis = self._torsion_scan(torsion, positions, r, theta, n_divisions=n_divisions)
xyzs = xyzs.value_in_unit_system(units.md_unit_system)
positions = positions.value_in_unit_system(units.md_unit_system)
for i, xyz in enumerate(xyzs):
positions[atom_idx,:] = xyz
position_set = time.time()
growth_context.setPositions(positions)
position_time = time.time() - position_set
self._position_set_time += position_time
energy_computation_init = time.time()
state = growth_context.getState(getEnergy=True)
potential_energy = state.getPotentialEnergy()
energy_computation_time = time.time() - energy_computation_init
self._energy_time += energy_computation_time
logq_i = -beta*potential_energy
logq[i] = logq_i
if np.sum(np.isnan(logq)) == n_divisions:
raise Exception("All %d torsion energies in torsion PMF are NaN." % n_divisions)
logq[np.isnan(logq)] = -np.inf
logq -= max(logq)
q = np.exp(logq)
Z = np.sum(q)
logp_torsions = logq - np.log(Z)
if hasattr(self, '_proposal_pdbfile'):
# Write proposal probabilities to PDB file as B-factors for inert atoms
f_i = -logp_torsions
f_i -= f_i.min() # minimum free energy is zero
f_i[f_i > 999.99] = 999.99
self._proposal_pdbfile.write('MODEL\n')
for i, xyz in enumerate(xyzs):
self._proposal_pdbfile.write('ATOM %5d %4s %3s %c%4d %8.3f%8.3f%8.3f%6.2f%6.2f\n' % (i+1, ' Ar ', 'Ar ', ' ', atom_idx+1, 10*xyz[0], 10*xyz[1], 10*xyz[2], np.exp(logp_torsions[i]), f_i[i]))
self._proposal_pdbfile.write('TER\n')
self._proposal_pdbfile.write('ENDMDL\n')
# TODO: Write proposal PMFs to storage
# atom_proposal_indices[order]
# atom_positions[order,k]
# torsion_pmf[order, division_index]
return logp_torsions, phis
def _propose_torsion(self, growth_context, torsion, positions, r, theta, beta, n_divisions=360):
"""
Propose a torsion using OpenMM
Parameters
----------
growth_context : openmm.Context
Context containing the modified system and
torsion : parmed.Dihedral
parmed Dihedral containing relevant atoms
positions : [n,3] np.ndarray in nm
positions of the atoms in the system
r : float in nm
bond length
theta : float in radians
bond angle
beta : float
inverse temperature
n_divisions : int, optional
number of divisions for the torsion scan. default 360
Returns
-------
phi : float in radians
The proposed torsion
logp : float
The log probability of the proposal.
"""
logp_torsions, phis = self._torsion_log_probability_mass_function(growth_context, torsion, positions, r, theta, beta, n_divisions=n_divisions)
division = units.Quantity(2*np.pi/n_divisions, unit=units.radian)
phi_median_idx = np.random.choice(range(len(phis)), p=np.exp(logp_torsions))
phi_min = phis[phi_median_idx] - division/2.0
phi_max = phis[phi_median_idx] + division/2.0
phi = np.random.uniform(phi_min.value_in_unit(units.radian), phi_max.value_in_unit(units.radian))
logp = logp_torsions[phi_median_idx] - np.log(2*np.pi / n_divisions) # convert from probability mass function to probability density function so that sum(dphi*p) = 1, with dphi = (2*pi)/n_divisions
return units.Quantity(phi, unit=units.radian), logp
def _torsion_logp(self, growth_context, torsion, positions, r, theta, phi, beta, n_divisions=360):
"""
Calculate the logp of a torsion using OpenMM
Parameters
----------
growth_context : openmm.Context
Context containing the modified system and
torsion : parmed.Dihedral
parmed Dihedral containing relevant atoms
positions : [n,3] np.ndarray in nm
positions of the atoms in the system
r : float in nm
Bond length
theta : float in radians
Bond angle
phi : float, in radians
The torsion angle
beta : float
inverse temperature
n_divisions : int, optional
number of divisions for logp calculation. default 360.
Returns
-------
torsion_logp : float
the logp of this torsion
"""
logp_torsions, phis = self._torsion_log_probability_mass_function(growth_context, torsion, positions, r, theta, beta, n_divisions=n_divisions)
phi_idx = np.argmin(np.abs(phi-phis)) # WARNING: This assumes both phi and phis have domain of [-pi,+pi)
torsion_logp = logp_torsions[phi_idx] - np.log(2*np.pi / n_divisions) # convert from probability mass function to probability density function so that sum(dphi*p) = 1, with dphi = (2*pi)/n_divisions.
return torsion_logp
[docs]class PredAtomTopologyIndex(oechem.OEUnaryAtomPred):
[docs] def __init__(self, topology_index):
super(PredAtomTopologyIndex, self).__init__()
self._topology_index = topology_index
def __call__(self, atom):
atom_data = atom.GetData()
if 'topology_index' in atom_data.keys():
if atom_data['topology_index'] == self._topology_index:
return True
return False
[docs]class BootstrapParticleFilter(object):
"""
Implements a Bootstrap Particle Filter (BPF)
to sample from the appropriate degrees of freedom.
Designed for use with the dimension-matching scheme
of Perses.
"""
[docs] def __init__(self, growth_context, atom_torsions, initial_positions, beta, n_particles=18, resample_frequency=10):
"""
Parameters
----------
growth_context : simtk.openmm.Context object
Context containing appropriate "growth system"
atom_torsions : dict
parmed.Atom : parmed.Dihedral dict that specifies
what torsion to use to propose each atom
initial_positions : np.ndarray [n,3]
The positions of existing atoms.
beta : simtk.unit.Quantity
The inverse temperature, with units
n_particles : int, optional
The number of particles in the BPF (note that this
is NOT the number of atoms). Default 18.
resample_frequency : int, optional
How often to resample particles. default 10
"""
raise NotImplementedError("The implementation of this GeometryEngine is not complete")
self._system = growth_context.getSystem()
self._beta = beta
self._growth_stage = 0
self._growth_context = growth_context
self._atom_torsions = atom_torsions
self._n_particles = n_particles
self._resample_frequency = resample_frequency
self._n_new_atoms = len(self._atom_torsions)
self._initial_positions = initial_positions
self._new_indices = [atom.idx for atom in self._atom_torsions.keys()]
#create a matrix for log weights (n_particles, n_stages)
self._Wij = np.zeros([self._n_particles, self._n_new_atoms])
#create an array for positions--only store new positions to avoid
#consuming way too much memory
self._new_positions = np.zeros([self._n_particles, self._n_new_atoms, 3])
self._generate_configurations()
def _internal_to_cartesian(self, bond_position, angle_position, torsion_position, r, theta, phi):
"""
Calculate the cartesian coordinates given the internal, as well as abs(detJ)
"""
from perses.rjmc import coordinate_numba
r = r.value_in_unit(units.nanometers)
theta = theta.value_in_unit(units.radians)
phi = phi.value_in_unit(units.radians)
bond_position = bond_position.astype(np.float64)
angle_position = angle_position.astype(np.float64)
torsion_position = torsion_position.astype(np.float64)
xyz = coordinate_numba.internal_to_cartesian(bond_position, angle_position, torsion_position, np.array([r, theta, phi], dtype=np.float64))
return xyz, r**2*np.sin(theta)
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 _log_unnormalized_target(self, new_positions):
"""
Given a set of new positions (not all positions!) and a growth
stage, return the log unnormalized probability.
Parameters
----------
new_positions : np.array
Array containing m 3D coordinates of new atoms
Returns
-------
log_unnormalized_probability : float
The unnormalized probability of this configuration
"""
positions = copy.deepcopy(self._initial_positions)
positions[self._new_indices] = new_positions
self._growth_context.setParameter('growth_stage', self._growth_stage)
self._growth_context.setPositions(positions)
energy = self._growth_context.getState(getEnergy=True).getPotentialEnergy()
return -self._beta*energy
def _get_relevant_angle(self, atom1, atom2, atom3):
"""
Get the angle containing the 3 given atoms
"""
atom1_angles = set(atom1.angles)
atom2_angles = set(atom2.angles)
atom3_angles = set(atom3.angles)
relevant_angle_set = atom1_angles.intersection(atom2_angles, atom3_angles)
relevant_angle = relevant_angle_set.pop()
if type(relevant_angle.type.k) != units.Quantity:
relevant_angle_with_units = self._add_angle_units(relevant_angle)
else:
relevant_angle_with_units = relevant_angle
return relevant_angle_with_units
def _add_bond_units(self, bond):
"""
Add the correct units to a harmonic bond
Arguments
---------
bond : parmed bond object
The bond to get units
Returns
-------
"""
if type(bond.type.k)==units.Quantity:
return bond
bond.type.req = units.Quantity(bond.type.req, unit=units.angstrom)
bond.type.k = units.Quantity(2.0*bond.type.k, unit=units.kilocalorie_per_mole/units.angstrom**2)
return bond
def _add_angle_units(self, angle):
"""
Add the correct units to a harmonic angle
Arguments
----------
angle : parmed angle object
the angle to get unit-ed
Returns
-------
angle_with_units : parmed angle
The angle, but with units on its parameters
"""
if type(angle.type.k)==units.Quantity:
return angle
angle.type.theteq = units.Quantity(angle.type.theteq, unit=units.degree)
angle.type.k = units.Quantity(2.0*angle.type.k, unit=units.kilocalorie_per_mole/units.radian**2)
return angle
def _get_relevant_bond(self, atom1, atom2):
"""
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
---------
atom1 : parmed atom object
One of the atoms in the bond
atom2 : parmed.atom object
The other atom in the bond
Returns
-------
relevant_bond_with_units : parmed.Bond
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 _bond_logq(self, r, bond):
"""
Calculate the log-probability of a given bond at a given inverse temperature
Arguments
---------
r : float
bond length, in nanometers
r0 : float
equilibrium bond length, in nanometers
k_eq : float
Spring constant of bond
beta : simtk.unit.Quantity
1/kT or inverse temperature
"""
k_eq = bond.type.k
r0 = bond.type.req
logq = -self._beta*0.5*k_eq*(r-r0)**2
return logq
def _angle_logq(self, theta, angle):
"""
Calculate the log-probability of a given bond at a given inverse temperature
Arguments
---------
theta : float
bond angle, in randians
angle : parmed angle object
Bond angle object containing parameters
beta : simtk.unit.Quantity
1/kT or inverse temperature
"""
k_eq = angle.type.k
theta0 = angle.type.theteq
logq = -self._beta*k_eq*0.5*(theta-theta0)**2
return logq
def _propose_bond(self, bond):
"""
Bond length proposal
"""
r0 = bond.type.req
k = bond.type.k
sigma_r = units.sqrt(1.0/(self._beta*k))
r = sigma_r*np.random.randn() + r0
return r
def _propose_angle(self, angle):
"""
Bond angle proposal
"""
theta0 = angle.type.theteq
k = angle.type.k
sigma_theta = units.sqrt(1.0/(self._beta*k))
theta = sigma_theta*np.random.randn() + theta0
return theta
def _propose_atom(self, atom, torsion, new_positions):
"""
Propose a set of internal coordinates (r, theta, phi) and transform
to cartesian coordinates (with jacobian correction).
for the given atom. R and theta are drawn from their respective
equilibrium distributions, whereas phi is simply a uniform sample.
Parameters
----------
atom : parmed.Atom
atom that will have its position proposed
torsion : parmed.Dihedral
torsion that contains relevant information for atom
new_positions : [m, 3] np.array
array of just the new positions (not existing atoms)
Returns
-------
xyz : [1,3] np.array of float
The proposed cartesian coordinates
logp : float
The log probability with jacobian correction
"""
positions = copy.deepcopy(self._initial_positions)
positions[self._new_indices] = new_positions
bond_atom = torsion.atom2
angle_atom = torsion.atom3
torsion_atom = torsion.atom4
if atom != torsion.atom1:
raise Exception('atom != torsion.atom1')
bond = self._get_relevant_bond(atom, bond_atom)
if bond is not None:
r = self._propose_bond(bond)
bond_k = bond.type.k
sigma_r = units.sqrt(1/(self._beta*bond_k))
logZ_r = np.log((np.sqrt(2*np.pi)*(sigma_r/units.angstroms))) # CHECK DOMAIN AND UNITS
logp_r = self._bond_logq(r, bond) - logZ_r
else:
constraint = self._get_bond_constraint(atom, bond_atom, self._system)
r = 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)
theta = self._propose_angle(angle)
angle_k = angle.type.k
sigma_theta = units.sqrt(1/(self._beta*angle_k))
logZ_theta = np.log((np.sqrt(2*np.pi)*(sigma_theta/units.radians))) # CHECK DOMAIN AND UNITS
logp_theta = self._angle_logq(theta, angle) - logZ_theta
#propose a torsion angle uniformly (this can be dramatically improved)
phi = np.random.uniform(-np.pi, np.pi)
logp_phi = -np.log(2*np.pi)
#get the new cartesian coordinates and detJ:
new_xyz, detJ = self._internal_to_cartesian(positions[bond_atom.idx], positions[angle_atom.idx], positions[torsion_atom.idx], r, theta, phi)
#accumulate logp
logp_proposal = logp_r + logp_theta + logp_phi + np.log(np.abs(detJ))
return new_xyz, logp_proposal
def _resample(self):
"""
Resample from the current set of weights and positions.
"""
particle_indices = range(self._n_particles)
new_indices = np.random.choice(particle_indices, size=self._n_particles, p=self._Wij[:, self._growth_stage-1])
for particle_index in particle_indices:
self._new_positions[particle_index, :, :] = self._new_positions[new_indices[particle_index], :, :]
self._Wij[:, self._growth_stage-1] = -np.log(self._n_particles) #set particle weights to be equal
def _generate_configurations(self):
"""
Generate the ensemble of configurations of the new atoms, approximately
from p(x_new | x_common).
"""
for i, atom_torsion in enumerate(self._atom_torsions.items()):
self._growth_stage = i+1
for particle_index in range(self._n_particles):
proposed_xyz, logp_proposal = self._propose_atom(atom_torsion[0], atom_torsion[1])
self._new_positions[particle_index, i, :] = proposed_xyz
unnormalized_log_target = self._log_unnormalized_target(self._new_positions[particle_index, :,:])
if i > 0:
self._Wij = [particle_index, i] = (unnormalized_log_target - logp_proposal) + self._Wij[particle_index, i-1]
else:
self._Wij = [particle_index, i] = unnormalized_log_target - logp_proposal
sum_log_weights = np.sum(np.exp(self._Wij[:,i]))
self._Wij -= np.log(sum_log_weights)
if i % self._resample_frequency == 0 and i != 0:
self._resample()
[docs]class OmegaGeometryEngine(GeometryEngine):
"""
This class proposes new small molecule geometries based on a set of precomputed
omega geometries.
"""
[docs] def __init__(self, n_omega_references=1, proposal_sigma=1.0, metadata=None):
self._n_omega_references = n_omega_references
self._proposal_sigma = 1.0
self._reference_oemols = {}
self._metadata = metadata
raise NotImplementedError("This GeometryEngine is not currently supported.")
[docs] def propose(self, top_proposal, current_positions, beta):
"""
Propose positions of new atoms according to a selected omega geometry.
Parameters
----------
top_proposal : TopologyProposal object
TopologyProposal object generated by the proposal engine
current_positions : [n, 3] np.array of float
Positions of the current system
beta : float
inverse temperature
Returns
-------
new_positions : [m, 3] np.array of float
The positions of the new system
logp_propose : float
The log-probability of the proposal
"""
pass
[docs] def logp_reverse(self, top_proposal, new_coordinates, old_coordinates, beta):
"""
Parameters
----------
top_proposal
new_coordinates
old_coordinates
beta
Returns
-------
"""
pass
[docs]class GeometrySystemGenerator(object):
"""
This is an internal utility class that generates OpenMM systems
with only valence terms and special parameters to assist in
geometry proposals.
"""
_HarmonicBondForceEnergy = "select(step({}+0.1 - growth_idx), (K/2)*(r-r0)^2, 0);"
_HarmonicAngleForceEnergy = "select(step({}+0.1 - growth_idx), (K/2)*(theta-theta0)^2, 0);"
_PeriodicTorsionForceEnergy = "select(step({}+0.1 - growth_idx), k*(1+cos(periodicity*theta-phase)), 0);"
[docs] def __init__(self, reference_system, growth_indices, parameter_name, add_extra_torsions=True, add_extra_angles=True, reference_topology=None, use_sterics=False, force_names=None, force_parameters=None, verbose=True):
"""
Parameters
----------
reference_system : simtk.openmm.System object
The system containing the relevant forces and particles
growth_indices : list of atom
The order in which the atom indices will be proposed
parameter_name : str
The name of the global context parameter
add_extra_torsions : bool, optional
Whether to add additional torsions to keep rings flat. Default true.
force_names : list of str
A list of the names of forces that will be included in this system
force_parameters : dict
Options for the forces (e.g., NonbondedMethod : 'CutffNonPeriodic')
verbose : bool, optional, default=False
If True, will print verbose output.
"""
ONE_4PI_EPS0 = 138.935456 # OpenMM constant for Coulomb interactions (openmm/platforms/reference/include/SimTKOpenMMRealType.h) in OpenMM units
# TODO: Replace this with an import from simtk.openmm.constants once these constants are available there
default_growth_index = len(growth_indices) # default value of growth index to use in System that is returned
self.current_growth_index = default_growth_index
# Nonbonded sterics and electrostatics.
# TODO: Allow user to select whether electrostatics or sterics components are included in the nonbonded interaction energy.
self._nonbondedEnergy = "select(step({}+0.1 - growth_idx), U_sterics + U_electrostatics, 0);"
self._nonbondedEnergy += "growth_idx = max(growth_idx1, growth_idx2);"
# Sterics
self._nonbondedEnergy += "U_sterics = 4*epsilon*x*(x-1.0); x = (sigma/r)^6;"
self._nonbondedEnergy += "epsilon = sqrt(epsilon1*epsilon2); sigma = 0.5*(sigma1 + sigma2);"
# Electrostatics
self._nonbondedEnergy += "U_electrostatics = ONE_4PI_EPS0*charge1*charge2/r;"
self._nonbondedEnergy += "ONE_4PI_EPS0 = %f;" % ONE_4PI_EPS0
# Exceptions (always included)
self._nonbondedExceptionEnergy = "select(step({}+0.1 - growth_idx), U_exception, 0);"
self._nonbondedExceptionEnergy += "U_exception = ONE_4PI_EPS0*chargeprod/r + 4*epsilon*x*(x-1.0); x = (sigma/r)^6;"
self._nonbondedExceptionEnergy += "ONE_4PI_EPS0 = %f;" % ONE_4PI_EPS0
self.sterics_cutoff_distance = 9.0 * units.angstroms # cutoff for sterics
self.verbose = verbose
# Get list of particle indices for new and old atoms.
new_particle_indices = [ atom.idx for atom in growth_indices ]
old_particle_indices = [idx for idx in range(reference_system.getNumParticles()) if idx not in new_particle_indices]
reference_forces = {reference_system.getForce(index).__class__.__name__ : reference_system.getForce(index) for index in range(reference_system.getNumForces())}
growth_system = openmm.System()
#create the forces:
modified_bond_force = openmm.CustomBondForce(self._HarmonicBondForceEnergy.format(parameter_name))
modified_bond_force.addPerBondParameter("r0")
modified_bond_force.addPerBondParameter("K")
modified_bond_force.addPerBondParameter("growth_idx")
modified_bond_force.addGlobalParameter(parameter_name, default_growth_index)
modified_angle_force = openmm.CustomAngleForce(self._HarmonicAngleForceEnergy.format(parameter_name))
modified_angle_force.addPerAngleParameter("theta0")
modified_angle_force.addPerAngleParameter("K")
modified_angle_force.addPerAngleParameter("growth_idx")
modified_angle_force.addGlobalParameter(parameter_name, default_growth_index)
modified_torsion_force = openmm.CustomTorsionForce(self._PeriodicTorsionForceEnergy.format(parameter_name))
modified_torsion_force.addPerTorsionParameter("periodicity")
modified_torsion_force.addPerTorsionParameter("phase")
modified_torsion_force.addPerTorsionParameter("k")
modified_torsion_force.addPerTorsionParameter("growth_idx")
modified_torsion_force.addGlobalParameter(parameter_name, default_growth_index)
growth_system.addForce(modified_bond_force)
growth_system.addForce(modified_angle_force)
growth_system.addForce(modified_torsion_force)
#copy the particles over
for i in range(reference_system.getNumParticles()):
growth_system.addParticle(reference_system.getParticleMass(i))
#copy each bond, adding the per-particle parameter as well
reference_bond_force = reference_forces['HarmonicBondForce']
for bond in range(reference_bond_force.getNumBonds()):
bond_parameters = reference_bond_force.getBondParameters(bond)
growth_idx = self._calculate_growth_idx(bond_parameters[:2], growth_indices)
if growth_idx==0:
continue
modified_bond_force.addBond(bond_parameters[0], bond_parameters[1], [bond_parameters[2], bond_parameters[3], growth_idx])
#copy each angle, adding the per particle parameter as well
reference_angle_force = reference_forces['HarmonicAngleForce']
for angle in range(reference_angle_force.getNumAngles()):
angle_parameters = reference_angle_force.getAngleParameters(angle)
growth_idx = self._calculate_growth_idx(angle_parameters[:3], growth_indices)
if growth_idx==0:
continue
modified_angle_force.addAngle(angle_parameters[0], angle_parameters[1], angle_parameters[2], [angle_parameters[3], angle_parameters[4], growth_idx])
#copy each torsion, adding the per particle parameter as well
reference_torsion_force = reference_forces['PeriodicTorsionForce']
for torsion in range(reference_torsion_force.getNumTorsions()):
torsion_parameters = reference_torsion_force.getTorsionParameters(torsion)
growth_idx = self._calculate_growth_idx(torsion_parameters[:4], growth_indices)
if growth_idx==0:
continue
modified_torsion_force.addTorsion(torsion_parameters[0], torsion_parameters[1], torsion_parameters[2], torsion_parameters[3], [torsion_parameters[4], torsion_parameters[5], torsion_parameters[6], growth_idx])
# Add (1,4) exceptions, regardless of whether 'use_sterics' is specified, because these are part of the valence forces.
if 'NonbondedForce' in reference_forces.keys():
custom_bond_force = openmm.CustomBondForce(self._nonbondedExceptionEnergy.format(parameter_name))
custom_bond_force.addPerBondParameter("chargeprod")
custom_bond_force.addPerBondParameter("sigma")
custom_bond_force.addPerBondParameter("epsilon")
custom_bond_force.addPerBondParameter("growth_idx")
custom_bond_force.addGlobalParameter(parameter_name, default_growth_index)
growth_system.addForce(custom_bond_force)
# Add exclusions, which are active at all times.
# (1,4) exceptions are always included, since they are part of the valence terms.
#print('growth_indices:', growth_indices)
reference_nonbonded_force = reference_forces['NonbondedForce']
for exception_index in range(reference_nonbonded_force.getNumExceptions()):
[particle_index_1, particle_index_2, chargeprod, sigma, epsilon] = reference_nonbonded_force.getExceptionParameters(exception_index)
growth_idx_1 = new_particle_indices.index(particle_index_1) + 1 if particle_index_1 in new_particle_indices else 0
growth_idx_2 = new_particle_indices.index(particle_index_2) + 1 if particle_index_2 in new_particle_indices else 0
growth_idx = max(growth_idx_1, growth_idx_2)
# Only need to add terms that are nonzero and involve newly added atoms.
if (growth_idx > 0) and ((chargeprod.value_in_unit_system(units.md_unit_system) != 0.0) or (epsilon.value_in_unit_system(units.md_unit_system) != 0.0)):
if self.verbose: _logger.info('Adding CustomBondForce: %5d %5d : chargeprod %8.3f e^2, sigma %8.3f A, epsilon %8.3f kcal/mol, growth_idx %5d' % (particle_index_1, particle_index_2, chargeprod/units.elementary_charge**2, sigma/units.angstrom, epsilon/units.kilocalorie_per_mole, growth_idx))
custom_bond_force.addBond(particle_index_1, particle_index_2, [chargeprod, sigma, epsilon, growth_idx])
#copy parameters for sterics parameters in nonbonded force
if 'NonbondedForce' in reference_forces.keys() and use_sterics:
modified_sterics_force = openmm.CustomNonbondedForce(self._nonbondedEnergy.format(parameter_name))
modified_sterics_force.addPerParticleParameter("charge")
modified_sterics_force.addPerParticleParameter("sigma")
modified_sterics_force.addPerParticleParameter("epsilon")
modified_sterics_force.addPerParticleParameter("growth_idx")
modified_sterics_force.addGlobalParameter(parameter_name, default_growth_index)
growth_system.addForce(modified_sterics_force)
# Translate nonbonded method to cutoff methods.
reference_nonbonded_force = reference_forces['NonbondedForce']
if reference_nonbonded_force in [openmm.NonbondedForce.NoCutoff, openmm.NonbondedForce.CutoffNonPeriodic]:
modified_sterics_force.setNonbondedMethod(openmm.CustomNonbondedForce.CutoffNonPeriodic)
elif reference_nonbonded_force in [openmm.NonbondedForce.CutoffPeriodic, openmm.NonbondedForce.PME, openmm.NonbondedForce.Ewald]:
modified_sterics_force.setNonbondedMethod(openmm.CustomNonbondedForce.CutoffPeriodic)
modified_sterics_force.setCutoffDistance(self.sterics_cutoff_distance)
# Add particle parameters.
for particle_index in range(reference_nonbonded_force.getNumParticles()):
[charge, sigma, epsilon] = reference_nonbonded_force.getParticleParameters(particle_index)
growth_idx = new_particle_indices.index(particle_index) + 1 if particle_index in new_particle_indices else 0
modified_sterics_force.addParticle([charge, sigma, epsilon, growth_idx])
if self.verbose and (growth_idx > 0):
_logger.info('Adding NonbondedForce particle %5d : charge %8.3f |e|, sigma %8.3f A, epsilon %8.3f kcal/mol, growth_idx %5d' % (particle_index, charge/units.elementary_charge, sigma/units.angstrom, epsilon/units.kilocalorie_per_mole, growth_idx))
# Add exclusions, which are active at all times.
# (1,4) exceptions are always included, since they are part of the valence terms.
for exception_index in range(reference_nonbonded_force.getNumExceptions()):
[particle_index_1, particle_index_2, chargeprod, sigma, epsilon] = reference_nonbonded_force.getExceptionParameters(exception_index)
modified_sterics_force.addExclusion(particle_index_1, particle_index_2)
# Only compute interactions of new particles with all other particles
# TODO: Allow inteactions to be resticted to only the residue being grown.
modified_sterics_force.addInteractionGroup(set(new_particle_indices), set(old_particle_indices))
modified_sterics_force.addInteractionGroup(set(new_particle_indices), set(new_particle_indices))
# Add extra ring-closing torsions, if requested.
if add_extra_torsions:
if reference_topology==None:
raise ValueError("Need to specify topology in order to add extra torsions.")
self._determine_extra_torsions(modified_torsion_force, reference_topology, growth_indices)
if add_extra_angles:
if reference_topology==None:
raise ValueError("Need to specify topology in order to add extra angles")
self._determine_extra_angles(modified_angle_force, reference_topology, growth_indices)
# Store growth system
self._growth_parameter_name = parameter_name
self._growth_system = growth_system
[docs] def set_growth_parameter_index(self, growth_parameter_index, context=None):
"""
Set the growth parameter index
"""
# TODO: Set default force global parameters if context is not None.
if context is not None:
context.setParameter(self._growth_parameter_name, growth_parameter_index)
self.current_growth_index = growth_parameter_index
[docs] def get_modified_system(self):
"""
Create a modified system with parameter_name parameter. When 0, only core atoms are interacting;
for each integer above 0, an additional atom is made interacting, with order determined by growth_index.
Returns
-------
growth_system : simtk.openmm.System object
System with the appropriate modifications, with growth parameter set to maximum.
"""
return self._growth_system
def _determine_extra_torsions(self, torsion_force, reference_topology, growth_indices):
"""
Determine which atoms need an extra torsion. First figure out which residue is
covered by the new atoms, then determine the rotatable bonds. Finally, construct
the residue in omega and measure the appropriate torsions, and generate relevant parameters.
ONLY ONE RESIDUE SHOULD BE CHANGING!
Parameters
----------
torsion_force : openmm.CustomTorsionForce object
the new/old torsion force if forward/backward
reference_topology : openmm.app.Topology object
the new/old topology if forward/backward
growth_indices : list of atom
The list of new atoms and the order in which they will be added.
Returns
-------
torsion_force : openmm.CustomTorsionForce
The torsion force with extra torsions added appropriately.
"""
# Do nothing if there are no atoms to grow.
if len(growth_indices) == 0:
return torsion_force
import openmoltools.forcefield_generators as forcefield_generators
atoms = list(reference_topology.atoms())
growth_indices = list(growth_indices)
#get residue from first atom
residue = atoms[growth_indices[0].idx].residue
try:
oemol = FFAllAngleGeometryEngine._oemol_from_residue(residue)
except Exception as e:
print("Could not generate an oemol from the residue.")
print(e)
# DEBUG: Write mol2 file.
debug = True
if debug:
if not hasattr(self, 'omega_index'):
self.omega_index = 0
filename = 'omega-%05d.mol2' % self.omega_index
print("Writing %s" % filename)
self.omega_index += 1
oemol_copy = oechem.OEMol(oemol)
ofs = oechem.oemolostream(filename)
oechem.OETriposAtomTypeNames(oemol_copy)
oechem.OEWriteMol2File(ofs, oemol_copy) # Preserve atom naming
ofs.close()
#get the omega geometry of the molecule:
omega = oeomega.OEOmega()
omega.SetMaxConfs(1)
omega.SetStrictStereo(False) #TODO: fix stereochem
omega(oemol)
#get the list of torsions in the molecule that are not about a rotatable bond
# Note that only torsions involving heavy atoms are enumerated here.
rotor = oechem.OEIsRotor()
torsion_predicate = oechem.OENotBond(rotor)
non_rotor_torsions = list(oechem.OEGetTorsions(oemol, torsion_predicate))
relevant_torsion_list = self._select_torsions_without_h(non_rotor_torsions)
#now, for each torsion, extract the set of indices and the angle
periodicity = 1
k = 120.0*units.kilocalories_per_mole # stddev of 12 degrees
#print([atom.name for atom in growth_indices])
for torsion in relevant_torsion_list:
#make sure to get the atom index that corresponds to the topology
atom_indices = [torsion.a.GetData("topology_index"), torsion.b.GetData("topology_index"), torsion.c.GetData("topology_index"), torsion.d.GetData("topology_index")]
# Determine phase in [-pi,+pi) interval
#phase = (np.pi)*units.radians+angle
phase = torsion.radians + np.pi # TODO: Check that this is the correct convention?
while (phase >= np.pi):
phase -= 2*np.pi
while (phase < -np.pi):
phase += 2*np.pi
phase *= units.radian
#print('PHASE>>>> ' + str(phase)) # DEBUG
growth_idx = self._calculate_growth_idx(atom_indices, growth_indices)
atom_names = [torsion.a.GetName(), torsion.b.GetName(), torsion.c.GetName(), torsion.d.GetName()]
#print("Adding torsion with atoms %s and growth index %d" %(str(atom_names), growth_idx))
#If this is a CustomTorsionForce, we need to pass the parameters as a list, and it will have the growth_idx parameter.
#If it's a regular PeriodicTorsionForce, there is no growth_index and the parameters are passed separately.
if isinstance(torsion_force, openmm.CustomTorsionForce):
torsion_force.addTorsion(atom_indices[0], atom_indices[1], atom_indices[2], atom_indices[3], [periodicity, phase, k, growth_idx])
elif isinstance(torsion_force, openmm.PeriodicTorsionForce):
torsion_force.addTorsion(atom_indices[0], atom_indices[1], atom_indices[2], atom_indices[3], periodicity, phase, k)
else:
raise ValueError("The force supplied to this method must be either a CustomTorsionForce or a PeriodicTorsionForce")
return torsion_force
def _select_torsions_without_h(self, torsion_list):
"""
Return only torsions that do not contain hydrogen
Parameters
----------
torsion_list : list of oechem.OETorsion
Returns
-------
heavy_torsions : list of oechem.OETorsion
"""
heavy_torsions = []
for torsion in torsion_list:
is_h_present = torsion.a.IsHydrogen() + torsion.b.IsHydrogen() + torsion.c.IsHydrogen() + torsion.d.IsHydrogen()
if not is_h_present:
heavy_torsions.append(torsion)
return heavy_torsions
def _determine_extra_angles(self, angle_force, reference_topology, growth_indices):
"""
Determine extra angles to be placed on aromatic ring members. Sometimes,
the native angle force is too weak to efficiently close the ring. As with the
torsion force, this method assumes that only one residue is changing at a time.
Parameters
----------
angle_force : simtk.openmm.CustomAngleForce
the force to which additional terms will be added
reference_topology : simtk.openmm.app.Topology
new/old topology if forward/backward
growth_indices : list of parmed.atom
Returns
-------
angle_force : simtk.openmm.CustomAngleForce
The modified angle force
"""
import itertools
if len(growth_indices)==0:
return
angle_force_constant = 400.0*units.kilojoules_per_mole/units.radians**2
atoms = list(reference_topology.atoms())
growth_indices = list(growth_indices)
#get residue from first atom
residue = atoms[growth_indices[0].idx].residue
try:
oemol = FFAllAngleGeometryEngine._oemol_from_residue(residue)
except Exception as e:
print("Could not generate an oemol from the residue.")
print(e)
#get the omega geometry of the molecule:
omega = oeomega.OEOmega()
omega.SetMaxConfs(1)
omega.SetStrictStereo(False) #TODO: fix stereochem
omega(oemol)
#we now have the residue as an oemol. Time to find the relevant angles.
#There's no equivalent to OEGetTorsions, so first find atoms that are relevant
#TODO: find out if that's really true
aromatic_pred = oechem.OEIsAromaticAtom()
heavy_pred = oechem.OEIsHeavy()
angle_criteria = oechem.OEAndAtom(aromatic_pred, heavy_pred)
#get all heavy aromatic atoms:
#TODO: do this more efficiently
heavy_aromatics = list(oemol.GetAtoms(angle_criteria))
for atom in heavy_aromatics:
#bonded_atoms = [bonded_atom for bonded_atom in list(atom.GetAtoms()) if bonded_atom in heavy_aromatics]
bonded_atoms = list(atom.GetAtoms())
for angle_atoms in itertools.combinations(bonded_atoms, 2):
angle = oechem.OEGetAngle(oemol, angle_atoms[0], atom, angle_atoms[1])
atom_indices = [angle_atoms[0].GetData("topology_index"), atom.GetData("topology_index"), angle_atoms[1].GetData("topology_index")]
angle_radians = angle*units.radian
growth_idx = self._calculate_growth_idx(atom_indices, growth_indices)
#If this is a CustomAngleForce, we need to pass the parameters as a list, and it will have the growth_idx parameter.
#If it's a regular HarmonicAngleForce, there is no growth_index and the parameters are passed separately.
if isinstance(angle_force, openmm.CustomAngleForce):
angle_force.addAngle(atom_indices[0], atom_indices[1], atom_indices[2], [angle_radians, angle_force_constant, growth_idx])
elif isinstance(angle_force, openmm.HarmonicAngleForce):
angle_force.addAngle(atom_indices[0], atom_indices[1], atom_indices[2], angle_radians, angle_force_constant)
else:
raise ValueError("Angle force must be either CustomAngleForce or HarmonicAngleForce")
return angle_force
def _calculate_growth_idx(self, particle_indices, growth_indices):
"""
Utility function to calculate the growth index of a particular force.
For each particle index, it will check to see if it is in growth_indices.
If not, 0 is added to an array, if yes, the index in growth_indices is added.
Finally, the method returns the max of the accumulated array
Parameters
----------
particle_indices : list of int
The indices of particles involved in this force
growth_indices : list of atom
The ordered list of indices for atom position proposals
Returns
-------
growth_idx : int
The growth_idx parameter
"""
growth_indices_list = [atom.idx for atom in list(growth_indices)]
particle_indices_set = set(particle_indices)
growth_indices_set = set(growth_indices_list)
new_atoms_in_force = particle_indices_set.intersection(growth_indices_set)
if len(new_atoms_in_force) == 0:
return 0
new_atom_growth_order = [growth_indices_list.index(atom_idx)+1 for atom_idx in new_atoms_in_force]
return max(new_atom_growth_order)
[docs]class GeometrySystemGeneratorFast(GeometrySystemGenerator):
"""
Use updateParametersInContext to make energy evaluation fast.
"""
[docs] def __init__(self, reference_system, growth_indices, parameter_name, add_extra_torsions=True, add_extra_angles=True, reference_topology=None, use_sterics=True, force_names=None, force_parameters=None, verbose=True):
"""
Parameters
----------
reference_system : simtk.openmm.System object
The system containing the relevant forces and particles
growth_indices : list of atom
The order in which the atom indices will be proposed
parameter_name : str
The name of the global context parameter
add_extra_torsions : bool, optional
Whether to add additional torsions to keep rings flat. Default true.
force_names : list of str
A list of the names of forces that will be included in this system
force_parameters : dict
Options for the forces (e.g., NonbondedMethod : 'CutffNonPeriodic')
verbose : bool, optional, default=False
If True, will print verbose output.
NB: We assume `reference_system` remains unmodified
"""
self.sterics_cutoff_distance = 9.0 * units.angstroms # cutoff for sterics
self.verbose = verbose
# Get list of particle indices for new and old atoms.
self._new_particle_indices = [ atom.idx for atom in growth_indices ]
self._old_particle_indices = [idx for idx in range(reference_system.getNumParticles()) if idx not in self._new_particle_indices]
self._growth_indices = growth_indices
# Determine forces to keep
forces_to_keep = ['HarmonicBondForce', 'HarmonicAngleForce', 'PeriodicTorsionForce']
if use_sterics:
forces_to_keep += ['NonbondedForce']
# Create reference system, removing forces we won't use
self._reference_system = copy.deepcopy(reference_system)
force_indices_to_remove = list()
for force_index in range(self._reference_system.getNumForces()):
force = self._reference_system.getForce(force_index)
force_name = force.__class__.__name__
if force_name not in forces_to_keep:
force_indices_to_remove.append(force_index)
for force_index in force_indices_to_remove[::-1]:
self._reference_system.removeForce(force_index)
# Create new system, copying forces we will keep.
self._growth_system = copy.deepcopy(self._reference_system)
#Extract the forces from the system to use for adding auxiliary angles and torsions
reference_forces = {reference_system.getForce(index).__class__.__name__ : reference_system.getForce(index) for index in range(reference_system.getNumForces())}
# Ensure 'canonical form' of System has all parameters turned on, or else we'll run into nonbonded exceptions
self.current_growth_index = -1
self.set_growth_parameter_index(len(self._growth_indices))
# Add extra ring-closing torsions, if requested.
if add_extra_torsions:
if reference_topology==None:
raise ValueError("Need to specify topology in order to add extra torsions.")
self._determine_extra_torsions(reference_forces['PeriodicTorsionForce'], reference_topology, growth_indices)
if add_extra_angles:
if reference_topology==None:
raise ValueError("Need to specify topology in order to add extra angles")
self._determine_extra_angles(reference_forces['HarmonicAngleForce'], reference_topology, growth_indices)
# TODO: Precompute growth indices for force terms for speed
[docs] def set_growth_parameter_index(self, growth_index, context=None):
"""
Set the growth parameter index
"""
self.current_growth_index = growth_index
for (growth_force, reference_force) in zip(self._growth_system.getForces(), self._reference_system.getForces()):
force_name = growth_force.__class__.__name__
if (force_name == 'HarmonicBondForce'):
for bond in range(reference_force.getNumBonds()):
parameters = reference_force.getBondParameters(bond)
this_growth_index = self._calculate_growth_idx(parameters[:2], self._growth_indices)
if (growth_index < this_growth_index):
parameters[3] *= 0.0
growth_force.setBondParameters(bond, *parameters)
elif (force_name == 'HarmonicAngleForce'):
for angle in range(reference_force.getNumAngles()):
parameters = reference_force.getAngleParameters(angle)
this_growth_index = self._calculate_growth_idx(parameters[:3], self._growth_indices)
if (growth_index < this_growth_index):
parameters[4] *= 0.0
growth_force.setAngleParameters(angle, *parameters)
elif (force_name == 'PeriodicTorsionForce'):
for torsion in range(reference_force.getNumTorsions()):
parameters = reference_force.getTorsionParameters(torsion)
this_growth_index = self._calculate_growth_idx(parameters[:4], self._growth_indices)
if (growth_index < this_growth_index):
parameters[6] *= 0.0
growth_force.setTorsionParameters(torsion, *parameters)
elif (force_name == 'NonbondedForce'):
for particle_index in range(reference_force.getNumParticles()):
parameters = reference_force.getParticleParameters(particle_index)
this_growth_index = self._calculate_growth_idx([particle_index], self._growth_indices)
if (growth_index < this_growth_index):
parameters[0] *= 0.0
parameters[2] *= 0.0
growth_force.setParticleParameters(particle_index, *parameters)
for exception_index in range(reference_force.getNumExceptions()):
parameters = reference_force.getExceptionParameters(exception_index)
this_growth_index = self._calculate_growth_idx(parameters[:2], self._growth_indices)
if (growth_index < this_growth_index):
parameters[2] *= 1.0e-6 # WORKAROUND // TODO: Change to zero when OpenMM issue is fixed
parameters[4] *= 1.0e-6 # WORKAROUND // TODO: Change to zero when OpenMM issue is fixed
growth_force.setExceptionParameters(exception_index, *parameters)
# Update parameters in context
if context is not None:
growth_force.updateParametersInContext(context)
[docs]class PredHBond(oechem.OEUnaryBondPred):
"""
Example elaborating usage on:
https://docs.eyesopen.com/toolkits/python/oechemtk/predicates.html#section-predicates-match
"""
def __call__(self, bond):
atom1 = bond.GetBgn()
atom2 = bond.GetEnd()
if atom1.IsHydrogen() or atom2.IsHydrogen():
return True
else:
return False
class NoTorsionError(Exception):
def __init__(self, message):
# Call the base class constructor with the parameters it needs
super(NoTorsionError, self).__init__(message)
