def set_nonbonded_method(
omm_sys: openmm.System,
key: str,
electrostatics: bool = True,
) -> openmm.System:
if key == "cutoff":
for force in omm_sys.getForces():
if type(force) == openmm.NonbondedForce:
force.setNonbondedMethod(openmm.NonbondedForce.CutoffPeriodic)
force.setCutoffDistance(0.9 * unit.nanometer)
force.setReactionFieldDielectric(1.0)
force.setUseDispersionCorrection(False)
force.setUseSwitchingFunction(False)
if not electrostatics:
for i in range(force.getNumParticles()):
params = force.getParticleParameters(i)
force.setParticleParameters(
i,
0,
params[1],
params[2],
)
elif key == "PME":
for force in omm_sys.getForces():
if type(force) == openmm.NonbondedForce:
force.setNonbondedMethod(openmm.NonbondedForce.PME)
force.setEwaldErrorTolerance(1e-6)
return omm_sys
def _get_charges_from_openmm_system(omm_sys: openmm.System):
for force in omm_sys.getForces():
if type(force) == openmm.NonbondedForce:
break
for idx in range(omm_sys.getNumParticles()):
param = force.getParticleParameters(idx)
yield param[0].value_in_unit(simtk_unit.elementary_charge)
def _get_force(openmm_sys: openmm.System, force_type):
forces = [f for f in openmm_sys.getForces() if type(f) == force_type]
if len(forces) > 1:
raise NotImplementedError(
"Not yet able to process duplicate forces types")
return forces[0]
def _get_lj_params_from_openmm_system(omm_sys: openmm.System):
for force in omm_sys.getForces():
if type(force) == openmm.NonbondedForce:
break
n_particles = omm_sys.getNumParticles()
sigmas = np.asarray([*_get_sigma_from_nonbonded_force(n_particles, force)])
epsilons = np.asarray([*_get_epsilon_from_nonbonded_force(n_particles, force)])
return sigmas, epsilons
def zero_dihedral_contribution(openmm_system: openmm.System, dihedral_indices: Tuple[int, int]):
"""
Author Simon Boothroyd
Link gist.github.com/SimonBoothroyd/667b50314c628aabe5064f0defb6ad8e
Zeroes out the contributions of a particular dihedral to the total potential
energy for a given OpenMM system object.
Parameters
----------
openmm_system:
The OpenMM system object which will be modified so the specified dihedral will
not contribute to the total potential energy of the system.
dihedral_indices:
The indices of the two central atoms in the dihedral whose contributions should
be zeroed.
"""
torsion_forces = [
force
for force in openmm_system.getForces()
if isinstance(force, openmm.PeriodicTorsionForce)
]
for torsion_force in torsion_forces:
for torsion_index in range(torsion_force.getNumTorsions()):
(
index_i,
index_j,
index_k,
index_l,
periodicity,
phase,
k,
) = torsion_force.getTorsionParameters(torsion_index)
if (
index_j not in [dihedral_indices[0], dihedral_indices[1]]
or index_k not in [dihedral_indices[0], dihedral_indices[1]]
):
continue
torsion_force.setTorsionParameters(
torsion_index,
index_i,
index_j,
index_k,
index_l,
periodicity,
phase,
0.0
)
def _get_openmm_energies(
omm_sys: openmm.System,
box_vectors,
positions,
round_positions=None,
hard_cutoff=False,
electrostatics: bool = True,
) -> EnergyReport:
if hard_cutoff:
omm_sys = set_nonbonded_method(omm_sys,
"cutoff",
electrostatics=electrostatics)
else:
omm_sys = set_nonbonded_method(omm_sys, "PME")
force_names = {force.__class__.__name__ for force in omm_sys.getForces()}
group_to_force = {
i: force_name
for i, force_name in enumerate(force_names)
}
force_to_group = {
force_name: i
for i, force_name in group_to_force.items()
}
for force in omm_sys.getForces():
force_name = force.__class__.__name__
force.setForceGroup(force_to_group[force_name])
integrator = openmm.VerletIntegrator(1.0 * unit.femtoseconds)
context = openmm.Context(omm_sys, integrator)
if not isinstance(box_vectors, unit.Quantity):
box_vectors = box_vectors.magnitude * unit.nanometer
context.setPeriodicBoxVectors(*box_vectors)
if isinstance(positions, unit.Quantity):
# Convert list of Vec3 into a NumPy array
positions = np.asarray(positions.value_in_unit(
unit.nanometer)) * unit.nanometer
else:
positions = positions.magnitude * unit.nanometer
if round_positions is not None:
rounded = np.round(positions, round_positions)
context.setPositions(rounded)
else:
context.setPositions(positions)
force_groups = {
force.getForceGroup()
for force in context.getSystem().getForces()
}
omm_energies = dict()
for force_group in force_groups:
state = context.getState(getEnergy=True, groups={force_group})
omm_energies[group_to_force[force_group]] = state.getPotentialEnergy()
del state
# Fill in missing keys if system does not have all typical forces
for required_key in [
"HarmonicBondForce",
"HarmonicAngleForce",
"PeriodicTorsionForce",
"NonbondedForce",
]:
if required_key not in omm_energies:
omm_energies[required_key] = 0.0 * kj_mol
del context
del integrator
report = EnergyReport()
report.energies.update({
"Bond": omm_energies["HarmonicBondForce"],
"Angle": omm_energies["HarmonicAngleForce"],
"Torsion": omm_energies["PeriodicTorsionForce"],
"Nonbonded": omm_energies["NonbondedForce"],
})
if "CustomNonbondedForce" in omm_energies:
report.energies["Nonbonded"] += omm_energies["CustomNonbondedForce"]
if "RBTorsionForce" in omm_energies:
report.energies["Torsion"] += omm_energies["RBTorsionForce"]
return report
def add_interactions(self, state: interfaces.IState, system: mm.System,
topology: app.Topology) -> mm.System:
# The approach we use is based on
# Habeck, Nilges, Rieping, J. Biomol. NMR., 2007, 135-144.
#
# Rather than solving for the exact alignment tensor
# every step, we sample from a distribution of alignment
# tensors.
#
# We encode the five components of the alignment tensor in
# the positions of two dummy atoms relative to the center
# of mass. The value of kappa should be scaled so that the
# components of the alignment tensor are approximately unity.
#
# There is a restraint on the z-component of the seocnd dummy
# particle to ensure that it does not diffuse off to ininity,
# which could cause precision issues.
if self.active:
rdc_force = mm.CustomCentroidBondForce(
5,
"Erest + z_scaler*Ez;"
"Erest = (1 - step(dev - quadcut)) * quad + step(dev - quadcut) * linear;"
"linear = 0.5 * k_rdc * quadcut^2 + k_rdc * quadcut * (dev - quadcut);"
"quad = 0.5 * k_rdc * dev^2;"
"dev = max(0, abs(d_obs - dcalc) - flat);"
"dcalc=2/3 * kappa_rdc/r^5 * (s1*(rx^2-ry^2) + s2*(3*rz^2-r^2) + s3*2*rx*ry + s4*2*rx*rz + s5*2*ry*rz);"
"r=distance(g4, g5);"
"rx=x4-x5;"
"ry=y4-y5;"
"rz=z4-z5;"
"s1=x2-x1;"
"s2=y2-y1;"
"s3=z2-z1;"
"s4=x3-x1;"
"s5=y3-y1;"
"Ez=(z3-z1)^2;",
)
rdc_force.addPerBondParameter("d_obs")
rdc_force.addPerBondParameter("kappa_rdc")
rdc_force.addPerBondParameter("k_rdc")
rdc_force.addPerBondParameter("flat")
rdc_force.addPerBondParameter("quadcut")
rdc_force.addPerBondParameter("z_scaler")
for experiment in self.expt_dict:
# find the set of all atoms involved in this experiment
com_ind = set()
for r in self.expt_dict[experiment]:
com_ind.add(r.atom_index_1)
com_ind.add(r.atom_index_2)
# add groups for the COM and dummy particles
s1 = self.expt_dict[experiment][0].s1_index
s2 = self.expt_dict[experiment][0].s2_index
g1 = rdc_force.addGroup(list(com_ind))
g2 = rdc_force.addGroup([s1])
g3 = rdc_force.addGroup([s2])
# add non-bonded exclusions between dummy particles and all other atoms
nb_forces = [
f for f in system.getForces()
if isinstance(f, mm.NonbondedForce)
or isinstance(f, mm.CustomNonbondedForce)
]
for nb_force in nb_forces:
n_parts = nb_force.getNumParticles()
for i in range(n_parts):
if isinstance(nb_force, mm.NonbondedForce):
if i != s1:
nb_force.addException(i,
s1,
0.0,
0.0,
0.0,
replace=True)
if i != s2:
nb_force.addException(i,
s2,
0.0,
0.0,
0.0,
replace=True)
else:
if i != s1:
nb_force.addExclusion(i, s1)
if i != s2:
nb_force.addExclusion(i, s2)
for r in self.expt_dict[experiment]:
# add groups for the atoms involved in the RDC
g4 = rdc_force.addGroup([r.atom_index_1])
g5 = rdc_force.addGroup([r.atom_index_2])
rdc_force.addBond(
[g1, g2, g3, g4, g5],
[
r.d_obs,
r.kappa,
0.0,
r.tolerance,
r.quadratic_cut,
0,
], # z_scaler initial value shouldn't matter
)
system.addForce(rdc_force)
self.force = rdc_force
return system
def _get_openmm_energies(
omm_sys: openmm.System,
box_vectors,
positions,
round_positions=None,
hard_cutoff=False,
electrostatics: bool = True,
) -> EnergyReport:
"""Given a prepared `openmm.System`, run a single-point energy calculation."""
"""\
if hard_cutoff:
omm_sys = _set_nonbonded_method(
omm_sys, "cutoff", electrostatics=electrostatics
)
else:
omm_sys = _set_nonbonded_method(omm_sys, "PME")
"""
for idx, force in enumerate(omm_sys.getForces()):
force.setForceGroup(idx)
integrator = openmm.VerletIntegrator(1.0 * unit.femtoseconds)
context = openmm.Context(omm_sys, integrator)
if box_vectors is not None:
if not isinstance(box_vectors, (unit.Quantity, list)):
box_vectors = box_vectors.magnitude * unit.nanometer
context.setPeriodicBoxVectors(*box_vectors)
if isinstance(positions, unit.Quantity):
# Convert list of Vec3 into a NumPy array
positions = np.asarray(positions.value_in_unit(
unit.nanometer)) * unit.nanometer
else:
positions = positions.magnitude * unit.nanometer
if round_positions is not None:
rounded = np.round(positions, round_positions)
context.setPositions(rounded)
else:
context.setPositions(positions)
raw_energies = dict()
omm_energies = dict()
for idx in range(omm_sys.getNumForces()):
state = context.getState(getEnergy=True, groups={idx})
raw_energies[idx] = state.getPotentialEnergy()
del state
# This assumes that only custom forces will have duplicate instances
for key in raw_energies:
force = omm_sys.getForce(key)
if type(force) == openmm.HarmonicBondForce:
omm_energies["HarmonicBondForce"] = raw_energies[key]
elif type(force) == openmm.HarmonicAngleForce:
omm_energies["HarmonicAngleForce"] = raw_energies[key]
elif type(force) == openmm.PeriodicTorsionForce:
omm_energies["PeriodicTorsionForce"] = raw_energies[key]
elif type(force) in [
openmm.NonbondedForce,
openmm.CustomNonbondedForce,
openmm.CustomBondForce,
]:
if "Nonbonded" in omm_energies:
omm_energies["Nonbonded"] += raw_energies[key]
else:
omm_energies["Nonbonded"] = raw_energies[key]
# Fill in missing keys if interchange does not have all typical forces
for required_key in [
"HarmonicBondForce",
"HarmonicAngleForce",
"PeriodicTorsionForce",
"NonbondedForce",
]:
if not any(required_key in val for val in omm_energies):
pass # omm_energies[required_key] = 0.0 * kj_mol
del context
del integrator
report = EnergyReport()
report.update_energies({
"Bond":
omm_energies.get("HarmonicBondForce", 0.0 * kj_mol),
"Angle":
omm_energies.get("HarmonicAngleForce", 0.0 * kj_mol),
"Torsion":
_canonicalize_torsion_energies(omm_energies),
"Nonbonded":
omm_energies.get("Nonbonded",
_canonicalize_nonbonded_energies(omm_energies)),
})
report.energies.pop("vdW")
report.energies.pop("Electrostatics")
return report
