def func(x):
xyz = independent_vars_to_xyz(x)
# these methods require 3d input
xyzlist = np.array([xyz])
my_bonds = core.bonds(xyzlist, ibonds).flatten()
my_angles = core.angles(xyzlist, iangles).flatten()
my_dihedrals = core.dihedrals(xyzlist, idihedrals).flatten()
d1 = 18.9*(my_bonds - bonds)
# Here's a hack to keep the bonds from becoming too short.
# print "\r %.4f %.4f" % (np.max(d1), np.max((my_bonds - bonds) / np.sqrt(my_bonds*bonds))),
d1 += 10.0*(my_bonds - bonds) / np.sqrt(my_bonds*bonds)
d2 = my_angles - angles
d3 = my_dihedrals - dihedrals
# d1 /= d1.shape[0]
# d2 /= d2.shape[0]
# d3 /= d3.shape[0]
# print max(d1), max((my_bonds - bonds) ** 2 / (my_bonds * bonds))
if xrefi != None:
d4 = (x - xrefi).flatten() * 18.9 * w_xref
error = np.r_[d1, d2, np.arctan2(np.sin(d3), np.cos(d3)), d4]
else:
error = np.r_[d1, d2, np.arctan2(np.sin(d3), np.cos(d3))]
if verbose:
print "RMS_ERROR", np.sqrt(np.mean(np.square(error)))
return error
def main():
from path_operations import union_connectivity
#np.random.seed(42)
xyzlist = 0.1*np.random.randn(7, 5, 3)
atom_names = ['C' for i in range(5)]
ibonds, iangles, idihedrals = union_connectivity(xyzlist, atom_names)
bonds = core.bonds(xyzlist, ibonds)
angles = core.angles(xyzlist, iangles)
dihedrals = core.dihedrals(xyzlist, idihedrals)
xyz_guess = xyzlist[0] + 0.025*np.random.rand(*xyzlist[0].shape)
x = least_squares_cartesian(bonds[0], ibonds, angles[0], iangles,
dihedrals[0], idihedrals, xyz_guess)
print x
def smooth_internal(xyzlist, atom_names, width, w_morse=0.0, **kwargs):
"""Smooth a trajectory by transforming to redundant, internal coordinates,
running a 1d timeseries smoothing algorithm on each DOF, and then
reconstructing a set of consistent cartesian coordinates.
TODO: write this function as a iterator that yields s_xyz, so that
they can be saved to disk (async) immediately when they're produced.
Parameters
----------
xyzlist : np.ndarray
Cartesian coordinates
atom_names : array_like of strings
The names of the atoms. Required for determing connectivity.
width : float
Width for the smoothing kernels
w_morse: float
Weight of the Morse potential in the smoothing
Other Parameters
----------------
bond_width : float
Override width just for the bond terms
angle_width : float
Override width just for the angle terms
dihedral_width : float
Override width just for the dihedral terms
xyzlist_guess :
Cartesian coordinates to use as a guess during the
reconstruction from internal
Returns
-------
smoothed_xyzlist : np.ndarray
"""
bond_width = kwargs.pop('bond_width', width)
angle_width = kwargs.pop('angle_width', width)
dihedral_width = kwargs.pop('dihedral_width', width)
xyzlist_guess = kwargs.pop('xyzlist_guess', xyzlist)
for key in kwargs.keys():
raise KeyError('Unrecognized key, %s' % key)
ibonds, iangles, idihedrals = None, None, None
s_bonds, s_angles, s_dihedrals = None, None, None
ibonds, iangles, idihedrals = union_connectivity(xyzlist, atom_names)
# get the internal coordinates in each frame
bonds = core.bonds(xyzlist, ibonds)
angles = core.angles(xyzlist, iangles)
dihedrals = core.dihedrals(xyzlist, idihedrals)
# run the smoothing
s_bonds = np.zeros_like(bonds)
s_angles = np.zeros_like(angles)
s_dihedrals = np.zeros_like(dihedrals)
for i in xrange(bonds.shape[1]):
#s_bonds[:, i] = buttersworth_smooth(bonds[:, i], width=bond_width)
s_bonds[:, i] = window_smooth(bonds[:, i],
window_len=bond_width,
window='hanning')
for i in xrange(angles.shape[1]):
#s_angles[:, i] = buttersworth_smooth(angles[:, i], width=angle_width)
s_angles[:, i] = window_smooth(angles[:, i],
window_len=angle_width,
window='hanning')
# filter the dihedrals with the angular smoother, that filters
# the sin and cos components separately
for i in xrange(dihedrals.shape[1]):
#s_dihedrals[:, i] = angular_smooth(dihedrals[:, i],
# smoothing_func=buttersworth_smooth, width=dihedral_width)
s_dihedrals[:, i] = angular_smooth(dihedrals[:, i],
smoothing_func=window_smooth,
window_len=dihedral_width,
window='hanning')
# compute the inversion for each frame
s_xyzlist = np.zeros_like(xyzlist_guess)
errors = np.zeros(len(xyzlist_guess))
# Thresholds for error and jitter
thre_jit = 3.0
w_xrefs = 0.0
for i, xyz_guess in enumerate(xyzlist_guess):
w_xref = 0.0
passed = False
corrected = False
while not passed:
passed = False
if i > 0:
xref = s_xyzlist[i - 1]
else:
xref = None
passed = True
r = least_squares_cartesian(
s_bonds[i],
ibonds,
s_angles[i],
iangles,
s_dihedrals[i],
idihedrals,
xyz_guess,
xref=xref,
w_xref=w_xref,
elem=atom_names,
w_morse=(0 if i in [0, len(xyzlist_guess) - 1] else w_morse))
s_xyzlist[i], errors[i] = r
if i > 0:
aligned0 = align_trajectory(
np.array([xyzlist[i], xyzlist[i - 1]]), 0)
aligned1 = align_trajectory(
np.array([s_xyzlist[i], s_xyzlist[i - 1]]), 0)
maxd0 = np.max(np.abs(aligned0[1] - aligned0[0]))
maxd1 = np.max(np.abs(aligned1[1] - aligned1[0]))
if maxd0 > 1e-5:
jit = maxd1 / maxd0
else:
jit = 0.0
else:
jit = 0.0
if (not passed) and jit < thre_jit:
passed = True
if w_xref >= 1.99:
w_xrefs = w_xref - 1.0
elif w_xref < 0.1:
w_xrefs = 0.0
else:
w_xrefs = w_xref / 1.5
elif not passed:
if w_xref == 0.0:
if w_xrefs > 0.0:
w_xref = w_xrefs
else:
w_xref = 2.0**10 / 3.0**10
else:
if w_xref >= 0.99:
w_xref += 1.0
else:
w_xref *= 1.5
print "jitter %f, trying anchor = %f\r" % (jit, w_xref),
corrected = True
# Print out a message if we had to correct it.
if corrected:
print '\rxyz: error %f max(dx) %f jitter %s anchor %f' % (
errors[i], maxd1, jit, w_xref)
if (i % 10) == 0:
print "\rWorking on frame %i / %i" % (i, len(xyzlist_guess)),
#return_value = (interweave(s_xyzlist), interweave(errors))
return_value = s_xyzlist, errors
return return_value
def fgrad(x, indicate = False):
""" Calculate the objective function and its derivatives. """
# If the optimization algorithm tries to calculate twice for the same point, do nothing.
# if x == fgrad.x0: return
xyz = independent_vars_to_xyz(x)
# these methods require 3d input
xyzlist = np.array([xyz])
my_bonds = core.bonds(xyzlist, ibonds).flatten()
my_angles = core.angles(xyzlist, iangles).flatten()
my_dihedrals = core.dihedrals(xyzlist, idihedrals, anchor=dihedrals).flatten()
# Deviations of internal coordinates from ideal values.
d1 = w1*(my_bonds - bonds)
d2 = w2*(my_angles - angles)
d3 = w3*(my_dihedrals - dihedrals)
# Include an optional term if we have an anchor point.
if xrefi != None:
d4 = (x - xrefi).flatten() * w1 * w_xref
fgrad.error = np.r_[d1, d2, np.arctan2(np.sin(d3), np.cos(d3)), d4]
else:
fgrad.error = np.r_[d1, d2, np.arctan2(np.sin(d3), np.cos(d3))]
# The objective function contains another contribution from the Morse potential.
fgrad.X = np.dot(fgrad.error, fgrad.error)
d1s = np.dot(d1, d1)
d2s = np.dot(d2, d2)
d3s = np.dot(d3, d3)
M = Molecule()
M.elem = elem
M.xyzs = [np.array(xyz)*10]
if w_morse != 0.0:
EMorse, GMorse = PairwiseMorse(M)
EMorse = EMorse[0]
GMorse = GMorse[0]
else:
EMorse = 0.0
GMorse = np.zeros((n_atoms, 3), dtype=float)
if indicate:
if fgrad.X0 != None:
print ("LSq: %.4f (%+.4f) Distance: %.4f (%+.4f) Angle: %.4f (%+.4f) Dihedral: %.4f (%+.4f) Morse: % .4f (%+.4f)" %
(fgrad.X, fgrad.X - fgrad.X0, d1s, d1s - fgrad.d1s0, d2s, d2s - fgrad.d2s0, d3s, d3s - fgrad.d3s0, EMorse, EMorse - fgrad.EM0)),
else:
print "LSq: %.4f Distance: %.4f Angle: %.4f Dihedral: %.4f Morse: % .4f" % (fgrad.X, d1s, d2s, d3s, EMorse),
fgrad.X0 = fgrad.X
fgrad.d1s0 = d1s
fgrad.d2s0 = d2s
fgrad.d3s0 = d3s
fgrad.EM0 = EMorse
fgrad.X += w_morse*EMorse
# Derivatives of internal coordinates w/r.t. Cartesian coordinates.
d_bonds = core.bond_derivs(xyz, ibonds) * w1
d_angles = core.angle_derivs(xyz, iangles) * w2
d_dihedrals = core.dihedral_derivs(xyz, idihedrals) * w3
if xrefi != None:
# the derivatives of the internal coordinates wrt the cartesian
# this is 2d, with shape equal to n_internal x n_cartesian
d_internal = np.vstack([gxyz_to_independent_vars(d_bonds.reshape((len(ibonds), -1))),
gxyz_to_independent_vars(d_angles.reshape((len(iangles), -1))),
gxyz_to_independent_vars(d_dihedrals.reshape((len(idihedrals), -1))),
np.eye(len(x)) * w1 * w_xref])
else:
# the derivatives of the internal coordinates wrt the cartesian
# this is 2d, with shape equal to n_internal x n_cartesian
d_internal = np.vstack([gxyz_to_independent_vars(d_bonds.reshape((len(ibonds), -1))),
gxyz_to_independent_vars(d_angles.reshape((len(iangles), -1))),
gxyz_to_independent_vars(d_dihedrals.reshape((len(idihedrals), -1)))])
# print fgrad.error.shape, d_internal.shape
# print d_internal.shape
# print xyz_to_independent_vars(d_internal).shape
fgrad.G = 2*np.dot(fgrad.error, d_internal)
fgrad.G += xyz_to_independent_vars(w_morse*GMorse.flatten())
def smooth_internal(xyzlist, atom_names, width, **kwargs):
"""Smooth a trajectory by transforming to redundant, internal coordinates,
running a 1d timeseries smoothing algorithm on each DOF, and then
reconstructing a set of consistent cartesian coordinates.
TODO: write this function as a iterator that yields s_xyz, so that
they can be saved to disk (async) immediately when they're produced.
Parameters
----------
xyzlist : np.ndarray
Cartesian coordinates
atom_names : array_like of strings
The names of the atoms. Required for determing connectivity.
width : float
Width for the smoothing kernels
Other Parameters
----------------
bond_width : float
Override width just for the bond terms
angle_width : float
Override width just for the angle terms
dihedral_width : float
Override width just for the dihedral terms
xyzlist_guess :
Cartesian coordinates to use as a guess during the
reconstruction from internal
Returns
-------
smoothed_xyzlist : np.ndarray
"""
bond_width = kwargs.pop('bond_width', width)
angle_width = kwargs.pop('angle_width', width)
dihedral_width = kwargs.pop('dihedral_width', width)
xyzlist_guess = kwargs.pop('xyzlist_guess', xyzlist)
for key in kwargs.keys():
raise KeyError('Unrecognized key, %s' % key)
ibonds, iangles, idihedrals = None, None, None
s_bonds, s_angles, s_dihedrals = None, None, None
with mpi_root():
ibonds, iangles, idihedrals = union_connectivity(xyzlist, atom_names)
# get the internal coordinates in each frame
bonds = core.bonds(xyzlist, ibonds)
angles = core.angles(xyzlist, iangles)
dihedrals = core.dihedrals(xyzlist, idihedrals)
# run the smoothing
s_bonds = np.zeros_like(bonds)
s_angles = np.zeros_like(angles)
s_dihedrals = np.zeros_like(dihedrals)
for i in xrange(bonds.shape[1]):
#s_bonds[:, i] = buttersworth_smooth(bonds[:, i], width=bond_width)
s_bonds[:, i] = window_smooth(bonds[:, i], window_len=bond_width, window='hanning')
for i in xrange(angles.shape[1]):
#s_angles[:, i] = buttersworth_smooth(angles[:, i], width=angle_width)
s_angles[:, i] = window_smooth(angles[:, i], window_len=angle_width, window='hanning')
# filter the dihedrals with the angular smoother, that filters
# the sin and cos components separately
for i in xrange(dihedrals.shape[1]):
#s_dihedrals[:, i] = angular_smooth(dihedrals[:, i],
# smoothing_func=buttersworth_smooth, width=dihedral_width)
s_dihedrals[:, i] = angular_smooth(dihedrals[:, i],
smoothing_func=window_smooth,
window_len=dihedral_width, window='hanning')
# group these into SIZE components, to be scattered
xyzlist_guess = group(xyzlist_guess, SIZE)
s_bonds = group(s_bonds, SIZE)
s_angles = group(s_angles, SIZE)
s_dihedrals = group(s_dihedrals, SIZE)
if RANK != 0:
xyzlist_guess = None
# scatter these
xyzlist_guess = COMM.scatter(xyzlist_guess, root=0)
s_bonds = COMM.scatter(s_bonds, root=0)
s_angles = COMM.scatter(s_angles, root=0)
s_dihedrals = COMM.scatter(s_dihedrals, root=0)
# broadcast the indices to every node
ibonds = COMM.bcast(ibonds, root=0)
iangles = COMM.bcast(iangles, root=0)
idihedrals = COMM.bcast(idihedrals, root=0)
# compute the inversion for each frame
s_xyzlist = np.zeros_like(xyzlist_guess)
errors = np.zeros(len(xyzlist_guess))
# Thresholds for error and jitter
thre_jit = 3.0
w_xrefs = 0.0
for i, xyz_guess in enumerate(xyzlist_guess):
w_xref = 0.0
passed = False
corrected = False
while not passed:
passed = False
if i > 0:
xref = s_xyzlist[i-1]
else:
xref = None
passed = True
r = least_squares_cartesian(s_bonds[i], ibonds, s_angles[i], iangles,
s_dihedrals[i], idihedrals, xyz_guess, xref=xref, w_xref=w_xref)
s_xyzlist[i], errors[i] = r
if i > 0:
aligned0 = align_trajectory(np.array([xyzlist[i],xyzlist[i-1]]), 0)
aligned1 = align_trajectory(np.array([s_xyzlist[i],s_xyzlist[i-1]]), 0)
maxd0 = np.max(np.abs(aligned0[1] - aligned0[0]))
maxd1 = np.max(np.abs(aligned1[1] - aligned1[0]))
if maxd0 > 1e-5:
jit = maxd1 / maxd0
else:
jit = 0.0
else:
jit = 0.0
if (not passed) and jit < thre_jit:
passed = True
if w_xref >= 1.99:
w_xrefs = w_xref - 1.0
elif w_xref < 0.1:
w_xrefs = 0.0
else:
w_xrefs = w_xref / 1.5
elif not passed:
if w_xref == 0.0:
if w_xrefs > 0.0:
w_xref = w_xrefs
else:
w_xref = 2.0**10 / 3.0**10
else:
if w_xref >= 0.99:
w_xref += 1.0
else:
w_xref *= 1.5
print "jitter %f, trying anchor = %f\r" % (jit, w_xref),
corrected = True
# Print out a message if we had to correct it.
if corrected:
print '\rRank %2d: (%3d)->xyz: error %f max(dx) %f jitter %s anchor %f' % (RANK, RANK + i*SIZE, errors[i], maxd1, jit, w_xref)
if (i%10) == 0:
print "\rWorking on frame %i / %i" % (i, len(xyzlist_guess)),
# gather the results back on root
s_xyzlist = COMM.gather(s_xyzlist, root=0)
errors = COMM.gather(errors, root=0)
return_value = (None, None)
with mpi_root():
# interleave the results back together
return_value = (interweave(s_xyzlist), interweave(errors))
return return_value