def _printEnergies(molecule, parameters, filename):
from ffevaluation.ffevaluate import FFEvaluate
energies = FFEvaluate(molecule,
parameters).calculateEnergies(molecule.coords[:, :,
0])
string = """
== Diagnostic Energies ==
Bond : {BOND_ENERGY:12.5f} kcal/mol
Angle : {ANGLE_ENERGY:12.5f} kcal/mol
Dihedral : {DIHEDRAL_ENERGY:12.5f} kcal/mol
Improper : {IMPROPER_ENERGY:12.5f} kcal/mol
Electro : {ELEC_ENERGY:12.5f} kcal/mol
VdW : {VDW_ENERGY:12.5f} kcal/mol
""".format(
BOND_ENERGY=energies["bond"],
ANGLE_ENERGY=energies["angle"],
DIHEDRAL_ENERGY=energies["dihedral"],
IMPROPER_ENERGY=energies["improper"],
ELEC_ENERGY=energies["elec"],
VDW_ENERGY=energies["vdw"],
)
for line in string.split("\n"):
logger.info(line)
with open(filename, "w") as file_:
file_.write(string)
logger.info("Write energy file: {}".format(filename))
def _evaluateConstTermsPerDihedral(self, parameters):
# Evaluate MM energies
const_energies = []
for idihed in range(self.numDihedrals):
# Zero the current dihedral values to calculate all other "constant" terms
modified_parameters = copy.deepcopy(parameters)
atomtypes = self._dihedral_atomtypes[idihed]
for term in modified_parameters.dihedral_types[atomtypes]:
term.phi_k = 0
assert term.per > 0 # Guard from messing up with improp
ff = FFEvaluate(self.molecule, modified_parameters)
scan_coords = self._coords[idihed]
energies = [
ff.calculateEnergies(coords[:, :, 0])["total"] for coords in scan_coords
]
const_energies.append(np.array(energies))
return const_energies
def _check(self):
# Evaluate the fitted energies
self._fitted_energies = []
ffeval = FFEvaluate(self.molecule, self.parameters)
for scan_coords in self._coords:
energies = [
ffeval.calculateEnergies(coords[:, :, 0])["total"]
for coords in scan_coords
]
self._fitted_energies.append(np.array(energies))
# Check the self-consistency of fitting
reference_energies = np.concatenate(self._reference_energies)
reference_energies -= np.mean(reference_energies)
fitted_energies = np.concatenate(self._fitted_energies)
fitted_energies -= np.mean(fitted_energies)
loss = np.sqrt(np.mean((fitted_energies - reference_energies) ** 2))
# HACK: without searches, the offset is not computed. So the test will not pass!
if self.num_iterations != 0:
assert np.isclose(self.loss, loss, rtol=0, atol=1e-2)
def _evaluateConstTerms(self):
"""
Evalutate constant MM terms
"""
parameters = copy.deepcopy(self.parameters)
# Disable parameterizable (i.e. non-constant) terms
for atomtypes in self._dihedral_atomtypes:
for term in parameters.dihedral_types[atomtypes]:
term.phi_k = 0
assert term.per > 0 # Guard from messing up with improper dihedrals
# Evaluate MM energies
const_energies = []
ff = FFEvaluate(self.molecule, parameters)
for scan_coords in self._coords:
energies = [
ff.calculateEnergies(coords[:, :, 0])["total"] for coords in scan_coords
]
const_energies.append(np.array(energies))
return const_energies
Epot = forces.compute(system.pos,
system.box,
system.forces,
returnDetails=True)[0]
myforces = system.forces.cpu().numpy()[0]
prm = parmed.charmm.CharmmParameterSet("./tests/water/parameters.prm")
struct = parmed.charmm.CharmmPsfFile("./tests/water/structure.psf")
keepForces(prm, struct, mol, forces=ommforceterm)
omm_energies, omm_forces = openmm_energy(prm,
struct,
mol.coords,
box=mol.box,
cutoff=7.3)
ffev = FFEvaluate(mol, prm, cutoff=7.3, rfa=True)
energy, forces_ffev, _ = ffev.calculate(mol.coords, mol.box)
energy = ffev.calculateEnergies(mol.coords, mol.box)
forces_ffev = forces_ffev.squeeze()
def compareForces(forces1, forces2):
return np.max(np.abs(forces1 - forces2).flatten())
datastore[forceterm[0]] = {
"omm": {
"energy": omm_energies["total"],
"forces": omm_forces
},
"torchmd": {
"energy": np.sum([x for _, x in Epot.items()]),
"forces": myforces
def _setup(self):
if len(self.dihedrals) != len(self.qm_results):
raise ValueError(
"The number of dihedral and QM result sets has to be the same!"
)
# Get dihedral names
self._names = [
"-".join(self.molecule.name[dihedral]) for dihedral in self.dihedrals
]
# Get all equivalent dihedrals
all_equivalent_dihedrals = detectParameterizableDihedrals(self.molecule)
all_equivalent_dihedrals = {
tuple(dihedrals[0]): dihedrals for dihedrals in all_equivalent_dihedrals
}
# Choose the selected dihedrals
equivalent_dihedrals = []
for dihedral, name in zip(self.dihedrals, self._names):
if tuple(dihedral) not in all_equivalent_dihedrals:
raise ValueError("{} is not a parameterizable dihedral!".format(name))
equivalent_dihedrals.append(all_equivalent_dihedrals[tuple(dihedral)])
# Get dihedral atom types
self._dihedral_atomtypes = [
findDihedralType(tuple(self.molecule.atomtype[dihedral]), self.parameters)
for dihedral in self.dihedrals
]
# Get reference QM energies and rotamer coordinates
self._reference_energies = []
self._coords = []
for results in self.qm_results:
self._reference_energies.append(
np.array([result.energy for result in results])
)
self._coords.append([result.coords for result in results])
# Calculate dihedral angle values
# [# of scans, # of dihedrals, # of conformations, # of equivalents]
self._angle_values = []
for scan_coords in self._coords:
scan_angle_values = []
for equivalent_indices in equivalent_dihedrals:
angle_values = []
for coords in scan_coords:
angle_values.append(
[
dihedralAngle(coords[indices, :, 0])
for indices in equivalent_indices
]
)
scan_angle_values.append(np.array(angle_values))
self._angle_values.append(scan_angle_values)
# Calculated initial MM energies
ff = FFEvaluate(self.molecule, self.parameters)
self._initial_energies = []
for scan_coords in self._coords:
energies = [
ff.calculateEnergies(coords[:, :, 0])["total"] for coords in scan_coords
]
self._initial_energies.append(np.array(energies))
# Make result directories
os.makedirs(self.result_directory, exist_ok=True)
def retrieve(self):
ff = FFEvaluate(self.molecule, self._parameters)
results = []
for iframe in range(self.molecule.numFrames):
self.molecule.frame = iframe
directory = os.path.join(self.directory, "%05d" % iframe)
os.makedirs(directory, exist_ok=True)
pickleFile = os.path.join(directory, "data.pkl")
if self._completed(directory):
with open(pickleFile, "rb") as fd:
result = pickle.load(fd)
logger.info("Loading QM data from %s" % pickleFile)
else:
result = QMResult()
result.errored = False
result.coords = self.molecule.coords[:, :,
iframe:iframe + 1].copy()
if self.optimize:
opt = nlopt.opt(nlopt.LN_COBYLA, result.coords.size)
opt.set_min_objective(lambda x, _: ff.calculateEnergies(
x.reshape((-1, 3)))["total"])
if self.restrained_dihedrals is not None:
for dihedral in self.restrained_dihedrals:
indices = dihedral.copy()
ref_angle = dihedralAngle(
self.molecule.coords[indices, :, iframe])
def constraint(x, _):
coords = x.reshape((-1, 3))
angle = dihedralAngle(coords[indices])
return np.sin(0.5 * (angle - ref_angle))
opt.add_equality_constraint(constraint)
opt.set_xtol_abs(1e-3) # Similar to Psi4 default
opt.set_maxeval(1000 * opt.get_dimension())
opt.set_initial_step(1e-3)
result.coords = opt.optimize(
result.coords.ravel()).reshape((-1, 3, 1))
logger.info("Optimization status: %d" %
opt.last_optimize_result())
result.energy = ff.calculateEnergies(result.coords[:, :,
0])["total"]
result.dipole = getDipole(self.molecule)
if self.optimize:
assert (opt.last_optimum_value() == result.energy
) # A self-consistency test
# Compute ESP values
if self.esp_points is not None:
assert self.molecule.numFrames == 1
result.esp_points = self.esp_points
distances = cdist(result.esp_points,
result.coords[:, :, 0]) # Angstrom
distances *= (const.physical_constants["Bohr radius"][0] /
const.angstrom) # Angstrom --> Bohr
result.esp_values = np.dot(
np.reciprocal(distances),
self.molecule.charge) # Hartree/Bohr
with open(pickleFile, "wb") as fd:
pickle.dump(result, fd)
results.append(result)
return results
def test_ffevaluate(self):
from natsort import natsorted
from ffevaluation.home import home
from moleculekit.molecule import Molecule
from glob import glob
import parmed
import os
import logging
logging.getLogger("parmed.structure").setLevel("ERROR")
for d in glob(os.path.join(home(dataDir="*"), "")):
with self.subTest(system=d):
print("\nRunning test:", d)
if os.path.basename(
os.path.abspath(d)) == "thrombin-ligand-amber":
abstol = 1e-1
elif os.path.basename(os.path.abspath(d)) == "waterbox":
abstol = 1e-3
else:
abstol = 1e-4
prmtopFile = glob(os.path.join(d, "*.prmtop"))
psfFile = glob(os.path.join(d, "*.psf"))
pdbFile = glob(os.path.join(d, "*.pdb"))
xtcFile = glob(os.path.join(d, "*.xtc"))
if len(glob(os.path.join(d, "*.prm"))):
prmFiles = [
fixParameters(glob(os.path.join(d, "*.prm"))[0]),
]
rtfFile = glob(os.path.join(d, "*.rtf"))
if len(rtfFile):
prmFiles.append(rtfFile[0])
else:
rtfFile = None
if len(psfFile):
mol = Molecule(psfFile[0])
elif len(prmtopFile):
mol = Molecule(prmtopFile[0])
if len(xtcFile):
mol.read(natsorted(xtcFile))
elif len(pdbFile):
mol.read(pdbFile[0])
else:
raise RuntimeError("No PDB or XTC")
coords = mol.coords
coords = coords[:, :, 0].squeeze()
rfa = False
cutoff = 0
if not np.all(mol.box == 0):
cutoff = np.min(mol.box) / 2 - 0.01
rfa = True
chargebackup = mol.charge.copy()
for force in (
"angle",
"bond",
"dihedral",
"lennardjones",
"electrostatic",
):
mol.charge = chargebackup.copy()
if len(psfFile):
struct = parmed.charmm.CharmmPsfFile(psfFile[0])
prm = parmed.charmm.CharmmParameterSet(*prmFiles)
keepForces(prm, struct, mol, forces=force)
elif len(prmtopFile):
struct = parmed.load_file(prmtopFile[0])
prm = parmed.amber.AmberParameterSet().from_structure(
struct)
keepForces(prm, struct, mol, forces=force)
keepForcesAmber(struct, mol, forces=force)
energies, forces, atmnrg = FFEvaluate(mol,
prm,
cutoff=cutoff,
rfa=rfa).calculate(
mol.coords,
mol.box)
energies = FFEvaluate.formatEnergies(energies[:, 0])
forces = forces[:, :, 0].squeeze()
omm_energies, omm_forces = openmm_energy(prm,
struct,
coords,
box=mol.box,
cutoff=cutoff)
ediff = compareEnergies(energies,
omm_energies,
abstol=abstol)
print(
" ",
force,
"Energy diff:",
ediff,
"Force diff:",
compareForces(forces, omm_forces),
)
if len(psfFile):
struct = parmed.charmm.CharmmPsfFile(psfFile[0])
prm = parmed.charmm.CharmmParameterSet(*prmFiles)
keepForces(prm, struct, mol)
elif len(prmtopFile):
struct = parmed.load_file(prmtopFile[0])
prm = parmed.amber.AmberParameterSet().from_structure(
struct)
keepForces(prm, struct, mol)
keepForcesAmber(struct, mol)
energies, forces, atmnrg = FFEvaluate(mol,
prm,
cutoff=cutoff,
rfa=rfa).calculate(
mol.coords, mol.box)
energies = FFEvaluate.formatEnergies(energies[:, 0])
forces = forces[:, :, 0].squeeze()
omm_energies, omm_forces = openmm_energy(prm,
struct,
coords,
box=mol.box,
cutoff=cutoff)
ediff = compareEnergies(energies, omm_energies, abstol=abstol)
print(
"All forces. Total energy:",
energies["total"],
"Energy diff:",
ediff,
"Force diff:",
compareForces(forces, omm_forces),
)
