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(.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 _fit(self):
from copy import deepcopy
# Save the initial parameters
vector = self._paramsToVector(self.parameters, self._parameterizable_dihedral_atomtypes)
if self.zeroed_parameters:
vector[:] = 0
# Evaluate the MM potential with this dihedral zeroed out
# The objective function will try to fit to the delta between
# the QM potential and this modified MM potential
for key in self._parameterizable_dihedral_atomtypes:
for term in self.parameters.dihedral_types[key]:
term.phi_k = 0
# Now evaluate the FF without the dihedral being fitted
self._target_energies = []
ff = FFEvaluate(self.molecule, self.parameters)
for rotamer_coords, ref_energies in zip(self._coords, self._reference_energies):
energies = ref_energies - np.array([ff.run(coords[:, :, 0])['total'] for coords in rotamer_coords])
energies -= np.min(energies)
self._target_energies.append(energies)
self._all_target_energies = np.concatenate(self._target_energies)
# Optimize the parameters
logger.info('Start parameter optimization')
# vector = self._optimize_CRS2_LM(vector) # TODO this should work better, but it doesn't
vector = self._optimize_random_search(vector)
logger.info('Final RMSD: %f kcal/mol' % self._objective(vector, None))
logger.info('Finished parameter optimization')
# Update the target dihedral with the optimized parameters
self._vectorToParams(self.parameters, self._parameterizable_dihedral_atomtypes, vector)
return self.loss
def _fit(self):
from copy import deepcopy
# Save the initial parameters
vector = self._paramsToVector(self.parameters, self._parameterizable_dihedral_atomtypes)
if self.zeroed_parameters:
vector[:] = 0
# Evaluate the MM potential with this dihedral zeroed out
# The objective function will try to fit to the delta between
# the QM potential and this modified MM potential
for key in self._parameterizable_dihedral_atomtypes:
for term in self.parameters.dihedral_types[key]:
term.phi_k = 0
# Now evaluate the FF without the dihedral being fitted
self._target_energies = []
ff = FFEvaluate(self.molecule, self.parameters)
for rotamer_coords, ref_energies in zip(self._coords, self._reference_energies):
energies = ref_energies - np.array([ff.calculateEnergies(coords[:, :, 0])['total'] for coords in rotamer_coords])
energies -= np.min(energies)
self._target_energies.append(energies)
self._all_target_energies = np.concatenate(self._target_energies)
# Optimize the parameters
logger.info('Start parameter optimization')
# vector = self._optimize_CRS2_LM(vector) # TODO this should work better, but it doesn't
vector = self._optimize_random_search(vector)
logger.info('Final RMSD: %f kcal/mol' % self._objective(vector, None))
logger.info('Finished parameter optimization')
# Update the target dihedral with the optimized parameters
self._vectorToParams(self.parameters, self._parameterizable_dihedral_atomtypes, vector)
return self.loss
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 printEnergies(molecule, parameters, filename):
from htmd.ffevaluation.ffevaluate import FFEvaluate
assert molecule.numFrames == 1
energies = FFEvaluate(molecule, parameters).run(molecule.coords[:, :, 0])
string = '''
== Diagnostic Energies ==
Bond : {BOND_ENERGY}
Angle : {ANGLE_ENERGY}
Dihedral : {DIHEDRAL_ENERGY}
Improper : {IMPROPER_ENERGY}
Electro : {ELEC_ENERGY}
VdW : {VDW_ENERGY}
'''.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'])
sys.stdout.write(string)
with open(filename, 'w') as file_:
file_.write(string)
def _printEnergies(molecule, parameters, filename):
from htmd.ffevaluation.ffevaluate import FFEvaluate
energies = FFEvaluate(molecule,
parameters).calculateEnergies(molecule.coords[:, :,
0])
string = '''
== Diagnostic Energies ==
Bond : {BOND_ENERGY:12.6g} kcal/mol
Angle : {ANGLE_ENERGY:12.6g} kcal/mol
Dihedral : {DIHEDRAL_ENERGY:12.6g} kcal/mol
Improper : {IMPROPER_ENERGY:12.6g} kcal/mol
Electro : {ELEC_ENERGY:12.6g} kcal/mol
VdW : {VDW_ENERGY:12.6g} 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 _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 equivalent dihedral atom indices
self._equivalent_indices = []
for idihed, dihedral in enumerate(self.dihedrals):
found = False
for parameterizableDihedral in self._parameterizable_dihedrals:
if np.all(list(parameterizableDihedral[0]) == dihedral):
self._equivalent_indices.append(parameterizableDihedral)
found = True
break
if not found:
raise ValueError('%s is not recognized as a parameterizable dihedral\n' % self._names[idihed])
# Get reference QM energies and rotamer coordinates
self._valid_qm_results = self._getValidQMResults()
self._reference_energies = []
self._coords = []
for results in self._valid_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 for the fitted equivalent dihedral
self._angle_values = []
for rotamer_coords, equivalent_indices in zip(self._coords, self._equivalent_indices):
angle_values = []
for coords in rotamer_coords:
angle_values.append([np.rad2deg(dihedralAngle(coords[indices, :, 0])) for indices in equivalent_indices])
self._angle_values.append(np.array(angle_values))
self._angle_values_rad = [np.deg2rad(angle_values)[:, :, None] for angle_values in self._angle_values]
self._parameterizable_dihedral_atomtypes = [tuple(self.molecule.atomtype[idx]) for idx in self.dihedrals]
# Calculated initial MM energies
ff = FFEvaluate(self.molecule, self.parameters)
self._initial_energies = []
for rotamer_coords in self._coords:
self._initial_energies.append(np.array([ff.calculateEnergies(coords[:, :, 0])['total'] for coords in rotamer_coords]))
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 equivalent dihedral atom indices
self._equivalent_indices = []
for idihed, dihedral in enumerate(self.dihedrals):
found = False
for parameterizableDihedral in self._parameterizable_dihedrals:
if np.all(list(parameterizableDihedral[0]) == dihedral):
self._equivalent_indices.append(parameterizableDihedral)
found = True
break
if not found:
raise ValueError('%s is not recognized as a parameterizable dihedral\n' % self._names[idihed])
# Get reference QM energies and rotamer coordinates
self._valid_qm_results = self._getValidQMResults()
self._reference_energies = []
self._coords = []
for results in self._valid_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 for the fitted equivalent dihedral
self._angle_values = []
for rotamer_coords, equivalent_indices in zip(self._coords, self._equivalent_indices):
angle_values = []
for coords in rotamer_coords:
angle_values.append([np.rad2deg(dihedralAngle(coords[indices, :, 0])) for indices in equivalent_indices])
self._angle_values.append(np.array(angle_values))
self._angle_values_rad = [np.deg2rad(angle_values)[:, :, None] for angle_values in self._angle_values]
self._parameterizable_dihedral_atomtypes = [tuple(self.molecule.atomtype[idx]) for idx in self.dihedrals]
# Calculated initial MM energies
ff = FFEvaluate(self.molecule, self.parameters)
self._initial_energies = []
for rotamer_coords in self._coords:
self._initial_energies.append(np.array([ff.run(coords[:, :, 0])['total'] for coords in rotamer_coords]))
def _check(self):
# Evaluate the fitted energies
self._fitted_energies = []
ffeval = FFEvaluate(self.molecule, self.parameters)
for rotamer_coords in self._coords:
self._fitted_energies.append(np.array([ffeval.calculateEnergies(coords[:, :, 0])['total'] for coords in rotamer_coords]))
# TODO make the self-consistency test numerically robust
#reference_energies = np.concatenate([energies - np.mean(energies) for energies in self._reference_energies])
#fitted_energies = np.concatenate([energies - np.mean(energies) for energies in self._fitted_energies])
#check_loss = np.sqrt(np.mean((fitted_energies - reference_energies)**2))
#assert np.isclose(self.loss, check_loss)
if self.result_directory:
os.makedirs(self.result_directory, exist_ok=True)
self.plotConformerEnergies()
for idihed in range(len(self.dihedrals)):
self.plotDihedralEnergies(idihed)
def _check(self):
# Evaluate the fitted energies
self._fitted_energies = []
ffeval = FFEvaluate(self.molecule, self.parameters)
for rotamer_coords in self._coords:
self._fitted_energies.append(np.array([ffeval.run(coords[:, :, 0])['total'] for coords in rotamer_coords]))
# TODO make the self-consistency test numerically robust
#reference_energies = np.concatenate([energies - np.mean(energies) for energies in self._reference_energies])
#fitted_energies = np.concatenate([energies - np.mean(energies) for energies in self._fitted_energies])
#check_loss = np.sqrt(np.mean((fitted_energies - reference_energies)**2))
#assert np.isclose(self.loss, check_loss)
if self.result_directory:
os.makedirs(self.result_directory, exist_ok=True)
self.plotConformerEnergies()
for idihed in range(len(self.dihedrals)):
self.plotDihedralEnergies(idihed)
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
mol.charge = chargebackup.copy()
if len(psfFile):
struct = parmed.charmm.CharmmPsfFile(psfFile[0])
prm = parmed.charmm.CharmmParameterSet(*prmFiles)
fromstruct = False
keepForces(prm, struct, mol, forces=force)
elif len(prmtopFile):
struct = parmed.load_file(prmtopFile[0])
prm = parmed.amber.AmberParameterSet().from_structure(struct)
fromstruct = True
keepForces(prm, struct, mol, forces=force)
keepForcesAmber(struct, mol, forces=force)
energies, forces, atmnrg = FFEvaluate(mol,
prm,
fromstruct=fromstruct,
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])
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)