def printEnergies(molecule, filename):
assert molecule.numFrames == 1
energies = FFEvaluate(molecule).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 _makeDihedralFittingSetFromQMResults(self, atoms, results):
# Extract the valid QM poses and energies from the QM result set
# Evaluate the MM on those poses
ffe = FFEvaluate(self)
ret = QMFittingSet()
ret.name = "%s-%s-%s-%s" % (self._rtf.names[atoms[0]], self._rtf.names[
atoms[1]], self._rtf.names[atoms[2]], self._rtf.names[atoms[3]])
completed = 0
qmin = 1.e100
for q in results:
if q.completed and not q.errored:
if q.energy < qmin:
qmin = q.energy
completed = 0
for q in results:
if q.completed and not q.errored:
if (
q.energy - qmin
) < 20.: # Only fit against QM points < 20 kcal above the minimum
mmeval = ffe.evaluate(q.coords)
if mmeval["vdw"] < 200:
completed += 1
phi = getPhi(q.coords, atoms)
ret.qm.append(q.energy - qmin)
ret.mm_original.append(mmeval['total'])
ret.coords.append(q.coords)
if phi > 180.:
phi -= 360.
ret.phi.append(phi)
else:
print(
"Omitting optimised pose for phi=%f (MM VDW too high)"
% phi)
else:
print(
"Omitting optimised QM pose (QM energy too high %f)" %
q.energy)
mmin = min(ret.mm_original)
# roughly align the qm with the mm
for q in range(completed):
ret.mm_original[q] = ret.mm_original[q] - mmin
ret.N = completed
if completed < 5:
raise RuntimeError(
"Fewer than 5 valid QM points. Not enough to fit!")
return ret
def _duplicateAtomType(self, atom_index):
"""Duplicate the type of the specified atom
Duplicated types are named: original_name + "x" + number, e.g. ca --> cax0
"""
# Get a type name
type_ = self._rtf.type_by_index[atom_index]
# if the type is already duplicated
if re.match(FFMolecule._ATOM_TYPE_REG_EX, type_):
return
# Create a new atom type name
i = 0
while ('%sx%d' % (type_, i)) in self._rtf.types:
i += 1
newtype = '%sx%d' % (type_, i)
logger.info('Create a new atom type %s from %s' % (newtype, type_))
# Duplicate the type in RTF
# TODO: move to RTF class
self._rtf.type_by_index[atom_index] = newtype
self._rtf.mass_by_type[newtype] = self._rtf.mass_by_type[type_]
self._rtf.types.append(newtype)
self._rtf.type_by_name[self._rtf.names[atom_index]] = newtype
self._rtf.type_by_index[atom_index] = newtype
self._rtf.typeindex_by_type[
newtype] = self._rtf.typeindex_by_type[type_] + 1000
self._rtf.element_by_type[newtype] = self._rtf.element_by_type[type_]
# Rename the atom types of the equivalent atoms
for index in self._equivalent_atoms[atom_index]:
if atom_index != index:
assert not re.match(FFMolecule._ATOM_TYPE_REG_EX,
self._rtf.type_by_index[index])
self._rtf.type_by_index[index] = newtype
self._rtf.type_by_name[self._rtf.names[index]] = newtype
# PRM parameters are duplicated during FF evaluation
# TODO: move to PRM class
FFEvaluate(self).run(self.coords[:, :, 0])
def _duplicateTypeOfAtom(self, aidx):
# This duplicates the type of the specified atom
# First get the type
atype = self._rtf.type_by_index[aidx]
# perhaps the type is already a duplicate? if so
# remove the duplicated suffix
atype = re.sub("x[0123456789]+$", "", atype)
i = 0
# make the new type name
while ("%sx%d" % (atype, i)) in self._rtf.types:
i += 1
newtype = "%sx%d" % (atype, i)
print("Creating new type %s from %s for atom %s" %
(newtype, atype, self._rtf.names[aidx]))
# duplicate the type in the fields RTF -- todo: move to a method in the RTF
self._rtf.type_by_index[aidx] = newtype
self._rtf.mass_by_type[newtype] = self._rtf.mass_by_type[atype]
self._rtf.types.append(newtype)
self._rtf.type_by_name[self._rtf.names[aidx]] = newtype
self._rtf.type_by_index[aidx] = newtype
self._rtf.typeindex_by_type[
newtype] = self._rtf.typeindex_by_type[atype] + 1000
self._rtf.element_by_type[newtype] = self._rtf.element_by_type[atype]
# # Now also reset the type of any atoms that share equivalency
for bidx in self._equivalent_atoms[aidx]:
if aidx != bidx:
if "x" in self._rtf.type_by_index[bidx]:
raise RuntimeError(
"Equivalent atom already has a duplicated type: {} {}".
format(bidx, self._rtf.type_by_index[bidx]))
self._rtf.type_by_index[bidx] = newtype
self._rtf.type_by_name[self._rtf.names[bidx]] = newtype
# the PRM parameters will be automatically duplicated by forcing an ff evaluation
ffe = FFEvaluate(self)
ffe.evaluate(self.coords)
def _fit(self):
# Save the initial parameters
vector = self._paramsToVector(self.parameters)
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 param in self.parameters:
param.k0 = 0
self.molecule._prm.updateDihedral(self.parameters)
# Now evaluate the FF without the dihedral being fitted
self._target_energies = []
ff = FFEvaluate(self.molecule)
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.parameters = self._vectorToParams(vector, self.parameters)
self.molecule._prm.updateDihedral(self.parameters)
return self.loss
def _check(self):
# Evaluate the fitted energies
self._fitted_energies = []
ffeval = FFEvaluate(self.molecule)
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 run(self):
ff = FFEvaluate(self.molecule)
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.run(x.reshape((-1, 3)))['total'])
if self.restrained_dihedrals is not None:
for dihedral in self.restrained_dihedrals:
indices = dihedral.copy()
ref_angle = np.deg2rad(
dihedralAngle(self.molecule.coords[indices, :,
iframe]))
def constraint(x, _):
coords = x.reshape((-1, 3))
angle = np.deg2rad(
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.run(result.coords[:, :, 0])['total']
result.dipole = self.molecule.getDipole()
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 fitSoftTorsion(self, angle, geomopt=True):
bkp_coords = self.coords.copy()
phi_to_fit = None
frozens = []
for d in self._soft_dihedrals:
if (d.atoms == angle).all():
phi_to_fit = d
frozens.append(d.atoms)
else:
if not geomopt:
frozens.append(d.atoms)
if not phi_to_fit:
raise ValueError("specified phi is not a recognised soft dihedral")
self._makeDihedralUnique(phi_to_fit)
atoms = phi_to_fit.atoms
equivs = phi_to_fit.equivalents
# Number of rotamers for each dihedral to compute
nrotamer = 36
# Create a copy of molecule with nrotamer frames
mol = self.copy()
for _ in range(nrotamer - 1):
mol.appendFrames(self)
assert mol.numFrames == nrotamer
# Set rotamer coordinates
angles = np.linspace(-np.pi, np.pi, num=nrotamer, endpoint=False)
for frame, angle in enumerate(angles):
mol.frame = frame
mol.setDihedral(atoms, angle, bonds=mol.bonds)
dirname = "dihedral-single-point"
if geomopt:
dirname = "dihedral-opt"
dih_name = "%s-%s-%s-%s" % (self.name[atoms[0]], self.name[atoms[1]],
self.name[atoms[2]], self.name[atoms[3]])
fitdir = os.path.join(self.outdir, dirname, dih_name,
self.output_directory_name())
try:
os.makedirs(fitdir, exist_ok=True)
except:
raise OSError(
'Directory {} could not be created. Check if you have permissions.'
.format(fitdir))
qmset = QMCalculation(mol,
charge=self.netcharge,
directory=fitdir,
frozen=frozens,
optimize=geomopt,
theory=self.theory,
solvent=self.solvent,
basis=self.basis,
execution=self.execution,
code=self.qmcode)
ret = self._makeDihedralFittingSetFromQMResults(atoms, qmset.results())
# Get the initial parameters of the dihedral we are going to fit
param = self._prm.dihedralParam(self._rtf.type_by_index[atoms[0]],
self._rtf.type_by_index[atoms[1]],
self._rtf.type_by_index[atoms[2]],
self._rtf.type_by_index[atoms[3]])
# Save these parameters as the best fit (fit to beat)
best_param = np.zeros((13))
for t in range(6):
best_param[t] = param[t].k0
best_param[t + 6] = param[t].phi0
best_param[12] = 0.
# Evalaute the mm potential with this dihedral zeroed out
# The objective function will try to fit to the delta between
# The QM potential and the this modified mm potential
for t in param:
t.k0 = t.phi0 = 0.
#t.e14 = 1. # Use whatever e14 has been inherited for the type
self._prm.updateDihedral(param)
ffeval = FFEvaluate(self)
# Now evaluate the ff without the dihedral being fitted
for t in range(ret.N):
mm_zeroed = ffeval.run(ret.coords[t][:, :, 0])['total']
ret.mm_delta.append(ret.qm[t] - mm_zeroed)
ret.mm_zeroed.append(mm_zeroed)
mmin1 = min(ret.mm_zeroed)
mmin2 = min(ret.mm_delta)
for t in range(ret.N):
ret.mm_zeroed[t] = ret.mm_zeroed[t] - mmin1
ret.mm_delta[t] = ret.mm_delta[t] - mmin2
self._fitDihedral_results = ret
self._fitDihedral_phi = param
# Now measure all of the soft dihedrals phis that are mapped to this dihedral
ret.phis = []
for iframe in range(ret.N):
ret.phis.append([ret.phi[iframe]])
for atoms in equivs:
angle = dihedralAngle(ret.coords[iframe][atoms, :, 0])
ret.phis[iframe].append(angle)
best_chisq = self._fitDihedral_objective(best_param)
bar = ProgressBar(64, description="Fitting")
for iframe in range(64):
(bounds, start) = self._fitDihedral_make_bounds(iframe)
xopt = optimize.minimize(self._fitDihedral_objective,
start,
method="L-BFGS-B",
bounds=bounds,
options={'disp': False})
chisq = self._fitDihedral_objective(xopt.x)
if (chisq < best_chisq):
best_chisq = chisq
best_param = xopt.x
bar.progress()
bar.stop()
# Update the target dihedral with the optimized parameters
for iframe in range(6):
param[iframe].k0 = best_param[0 + iframe]
param[iframe].phi0 = best_param[6 + iframe]
self._prm.updateDihedral(param)
param = self._prm.dihedralParam(self._rtf.type_by_index[atoms[0]],
self._rtf.type_by_index[atoms[1]],
self._rtf.type_by_index[atoms[2]],
self._rtf.type_by_index[atoms[3]])
# Finally evaluate the fitted potential
ffeval = FFEvaluate(self)
for t in range(ret.N):
ret.mm_fitted.append(ffeval.run(ret.coords[t][:, :, 0])['total'])
mmin = min(ret.mm_fitted)
chisq = 0.
for t in range(ret.N):
ret.mm_fitted[t] = ret.mm_fitted[t] - mmin
delta = ret.mm_fitted[t] - ret.qm[t]
chisq = chisq + (delta * delta)
ret.chisq = chisq
# TODO Score it
self.coords = bkp_coords
return ret
def fitSoftTorsion(self, phi, geomopt=True):
found = False
phi_to_fit = None
frozens = []
dih_index = 0
i = 0
bkp_coords = self.coords.copy()
for d in self._soft_dihedrals:
if (d.atoms == phi).all():
phi_to_fit = d
dih_index = i
frozens.append(d.atoms)
else:
if not geomopt:
frozens.append(d.atoms)
i += 1
if not phi_to_fit:
raise ValueError("specified phi is not a recognised soft dihedral")
self._makeDihedralUnique(phi_to_fit)
atoms = phi_to_fit.atoms
left = phi_to_fit.left
right = phi_to_fit.right
equivs = phi_to_fit.equivalents
step = 10 # degrees
nstep = int(360 / step)
cset = np.zeros((self.natoms, 3, nstep))
i = 0
for phi in range(-180, 180, step):
cset[:, :, i] = setPhi(self.coords[:, :, 0], atoms, left, right,
phi)
i += 1
mol = self.copy()
mol.coords = cset
dirname = "dihedral-single-point"
if geomopt:
dirname = "dihedral-opt"
dih_name = "%s-%s-%s-%s" % (self.name[atoms[0]], self.name[atoms[1]],
self.name[atoms[2]], self.name[atoms[3]])
fitdir = os.path.join(self.outdir, dirname, dih_name,
self.output_directory_name())
try:
os.makedirs(fitdir, exist_ok=True)
except:
raise OSError(
'Directory {} could not be created. Check if you have permissions.'
.format(fitdir))
qmset = QMCalculation(mol,
charge=self.netcharge,
directory=fitdir,
frozen=frozens,
optimize=geomopt,
theory=self.theory,
solvent=self.solvent,
basis=self.basis,
execution=self.execution,
code=self.qmcode)
ret = self._makeDihedralFittingSetFromQMResults(atoms, qmset.results())
# Get the initial parameters of the dihedral we are going to fit
param = self._prm.dihedralParam(self._rtf.type_by_index[atoms[0]],
self._rtf.type_by_index[atoms[1]],
self._rtf.type_by_index[atoms[2]],
self._rtf.type_by_index[atoms[3]])
# Save these parameters as the best fit (fit to beat)
best_param = np.zeros((13))
for t in range(6):
best_param[t] = param[t].k0
best_param[t + 6] = param[t].phi0
best_param[12] = 0.
# print(param)
# Evalaute the mm potential with this dihedral zeroed out
# The objective function will try to fit to the delta between
# The QM potential and the this modified mm potential
for t in param:
t.k0 = t.phi0 = 0.
#t.e14 = 1. # Use whatever e14 has been inherited for the type
self._prm.updateDihedral(param)
ffe = FFEvaluate(self)
# print(ffe.evaluate( ret.coords[0] ) )
# input
# Now evaluate the ff without the dihedral being fitted
for t in range(ret.N):
mm_zeroed = (ffe.evaluate(ret.coords[t])["total"])
ret.mm_delta.append(ret.qm[t] - mm_zeroed)
ret.mm_zeroed.append(mm_zeroed)
mmin1 = min(ret.mm_zeroed)
mmin2 = min(ret.mm_delta)
for t in range(ret.N):
ret.mm_zeroed[t] = ret.mm_zeroed[t] - mmin1
ret.mm_delta[t] = ret.mm_delta[t] - mmin2
self._fitDihedral_results = ret
self._fitDihedral_phi = param
# Now measure all of the soft dihedrals phis that are mapped to this dihedral
ret.phis = []
for i in range(ret.N):
ret.phis.append([ret.phi[i]])
for e in equivs:
ret.phis[i].append(getPhi(ret.coords[i], e))
# print ("EQUIVALENT DIHEDRALS FOR THIS DIHEDRAL" )
# print(equivs)
# print ("PHI VALUES TO FIT")
# print (ret.phis)
# Set up the NOLOPT fit
# There are 13 parameters, k,phi for n=1,2,3,4,5,6 and a shift
N = 13
# initial guess,
st = np.zeros(13)
# bounds
best_chisq = self._fitDihedral_objective(best_param)
# print("CHISQ of initial = %f" % ( best_chisq ) )
# Now zero out the terms of the dihedral we are going to fit
bar = ProgressBar(64, description="Fitting")
for i in range(64):
(bounds, start) = self._fitDihedral_make_bounds(i)
xopt = optimize.minimize(self._fitDihedral_objective,
start,
method="L-BFGS-B",
bounds=bounds,
options={'disp': False})
chisq = self._fitDihedral_objective(xopt.x)
# print( "CHISQ of fit = %f " % (chisq) )
if (chisq < best_chisq):
best_chisq = chisq
best_param = xopt.x
bar.progress()
bar.stop()
# print("Best ChiSQ = %f" %(best_chisq) )
# Update the target dihedral with the optimized parameters
# print(param)
# print(best_param )
for i in range(6):
param[i].k0 = best_param[0 + i]
param[i].phi0 = best_param[6 + i]
self._prm.updateDihedral(param)
# print(param)
param = self._prm.dihedralParam(self._rtf.type_by_index[atoms[0]],
self._rtf.type_by_index[atoms[1]],
self._rtf.type_by_index[atoms[2]],
self._rtf.type_by_index[atoms[3]])
# print(param)
# Finally evaluate the fitted potential
ffe = FFEvaluate(self)
for t in range(ret.N):
ret.mm_fitted.append(ffe.evaluate(ret.coords[t])["total"])
mmin = min(ret.mm_fitted)
chisq = 0.
# print( "QM energies" )
# print( ret.qm )
for t in range(ret.N):
ret.mm_fitted[t] = ret.mm_fitted[t] - mmin
delta = ret.mm_fitted[t] - ret.qm[t]
chisq = chisq + (delta * delta)
ret.chisq = chisq
# TODO Score it
self.coords = bkp_coords
return ret
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 = [
'%s-%s-%s-%s' % tuple(self.molecule.name[dihedral])
for dihedral in self.dihedrals
]
# Get equivalent dihedral atom indices
self._equivalent_indices = []
for idihed, dihedral in enumerate(self.dihedrals):
for rotatableDihedral in self.molecule._rotatable_dihedrals:
if np.all(rotatableDihedral.atoms == dihedral):
self._equivalent_indices.append(
[dihedral] + rotatableDihedral.equivalents)
break
else:
raise ValueError(
'%s is not recognized as a rotable dihedral\n' %
self._names[idihed])
# Get dihedral parameters
for dihedral in self.dihedrals:
self._makeDihedralUnique(dihedral)
all_types = [
tuple([
self.molecule._rtf.type_by_index[index] for index in dihedral
]) for dihedral in self.dihedrals
]
self.parameters = sum(
[self.molecule._prm.dihedralParam(*types) for types in all_types],
[])
# 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([
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
]
# Calculated initial MM energies
ff = FFEvaluate(self.molecule)
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
]))