def _find_nb_force(self, system: mm.System) -> None:
forces = [system.getForce(i) for i in range(system.getNumForces())]
nb_forces = [f for f in forces if isinstance(f, mm.NonbondedForce)]
if not nb_forces:
raise RuntimeError("REST2 could not find NonbondedForce")
if len(nb_forces) > 1:
raise RuntimeError("REST2 found more than one NonbondedForce")
self.nb_force = nb_forces[0]
def _find_dihedral_force(self, system: mm.System) -> None:
forces = [system.getForce(i) for i in range(system.getNumForces())]
dihed_forces = [
f for f in forces if isinstance(f, mm.PeriodicTorsionForce)
]
if not dihed_forces:
raise RuntimeError("REST2 could not find PeriodicTorsionForce")
if len(dihed_forces) > 1:
raise RuntimeError(
"REST2 found more than one PeriodicTorsionForce")
self.dihedral_force = dihed_forces[0]
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
def from_system(cls, system: openmm.System, indices: Dict[str, int]):
"""Instantiate a COOHDummyMover for a single C-COOH moiety in your system.
Parameters
----------
system - The OpenMM system containting the COOH moiety
indices - a dictionary labeling the indices of the C-COOH atoms with keys:
HO - index in the system of the hydroxyl hydrogen.
OH - index in the system of the hydroxyl oxygen
OC - index in the system of the carbonyl oxygen
CO - index in the system of the carbonyl carbon
R - index in the system of the atom "R" this COOH group is connected to.
"""
obj = cls()
# Hydroxyl hydrogen
obj.HO = indices["HO"]
# Hydroxyl oxygen
obj.OH = indices["OH"]
# Carbonyl oxygen
obj.OC = indices["OC"]
# The carbons are used as reference points for reflection
obj.CO = indices["CO"]
obj.R = indices["R"]
# All atoms that this class may decide to move
obj.movable = [indices["OC"], indices["OH"], indices["HO"]]
# The parameters for angles
obj.angles = []
# The parameters for dihedrals
obj.dihedrals = []
# Instantiate the class variable
# This is to keep track of angle and torsion force indices for all future instances
if (
COOHDummyMover.angleforceindex is None
or COOHDummyMover.dihedralforceindex is None
):
for force_index in range(system.getNumForces()):
force = system.getForce(force_index)
if force.__class__.__name__ == "HarmonicAngleForce":
COOHDummyMover.angleforceindex = force_index
elif force.__class__.__name__ == "PeriodicTorsionForce":
COOHDummyMover.dihedralforceindex = force_index
if COOHDummyMover.angleforceindex is None:
raise RuntimeError(
"{} requires the system to have a HarmonicAngleForce!".format(
COOHDummyMover.__name__
)
)
if COOHDummyMover.dihedralforceindex is None:
raise RuntimeError(
"{} requires the system to have a PeriodicTorsionForce!".format(
COOHDummyMover.__name__
)
)
angleforce = system.getForce(COOHDummyMover.angleforceindex)
torsionforce = system.getForce(COOHDummyMover.dihedralforceindex)
# Loop through and collect all angle energy terms that include moving atoms
for angle_index in range(angleforce.getNumAngles()):
*particles, theta0, k = angleforce.getAngleParameters(angle_index)
if any(particle in obj.movable for particle in particles):
# Energy function for this angle.
params = [k._value, theta0._value, *particles]
log.debug("Found this COOH angle: %s", params)
obj.angles.append(params)
# Loop through and collect all torsion energy terms that include moving atoms
for torsion_index in range(torsionforce.getNumTorsions()):
*particles, n, theta0, k = torsionforce.getTorsionParameters(torsion_index)
if any(particle in obj.movable for particle in particles):
# Energy function for this dihedral.
params = [k._value, n, theta0._value, *particles]
log.debug("Found this COOH dihedral: %s", params)
obj.dihedrals.append(params)
return obj