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
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
},
"ffeval": {
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