Created
January 30, 2013 03:51
-
-
Save rmcgibbo/4670468 to your computer and use it in GitHub Desktop.
Nudged Elastic Band in Python
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from __future__ import division | |
import numpy as np | |
import IPython as ip | |
import matplotlib.pyplot as pp | |
import warnings | |
from collections import namedtuple | |
import scipy.optimize | |
#symbolic algebra | |
import theano | |
import theano.tensor as T | |
Components = namedtuple('Components', ['parallel', 'perpindicular']) | |
def muller_potential(x, y): | |
"""Muller potential | |
Parameters | |
---------- | |
x : {float, np.ndarray, or theano symbolic variable} | |
X coordinate. If you supply an array, x and y need to be the same shape, | |
and the potential will be calculated at each (x,y pair) | |
y : {float, np.ndarray, or theano symbolic variable} | |
Y coordinate. If you supply an array, x and y need to be the same shape, | |
and the potential will be calculated at each (x,y pair) | |
Returns | |
------- | |
potential : {float, np.ndarray, or theano symbolic variable} | |
Potential energy. Will be the same shape as the inputs, x and y. | |
Reference | |
--------- | |
Code adapted from https://cims.nyu.edu/~eve2/ztsMueller.m | |
""" | |
aa = [-1, -1, -6.5, 0.7] | |
bb = [0, 0, 11, 0.6] | |
cc = [-10, -10, -6.5, 0.7] | |
AA = [-200, -100, -170, 15] | |
XX = [1, 0, -0.5, -1] | |
YY = [0, 0.5, 1.5, 1] | |
# use symbolic algebra if you supply symbolic quantities | |
exp = T.exp if isinstance(x, T.TensorVariable) else np.exp | |
value = 0 | |
for j in range(0, 4): | |
value += AA[j] * exp(aa[j] * (x - XX[j])**2 + \ | |
bb[j] * (x - XX[j]) * (y - YY[j]) + cc[j] * (y - YY[j])**2) | |
return value | |
def muller_grad_factory(): | |
"""Compile a theano function to compute the negative grad | |
of the muller potential""" | |
sym_x, sym_y = T.scalar(), T.scalar() | |
sym_V = muller_potential(sym_x, sym_y) | |
sym_F = T.grad(sym_V, [sym_x, sym_y]) | |
# F takes two arguments, x,y and returns a 2 | |
# element python list | |
F = theano.function([sym_x, sym_y], sym_F) | |
def force(position): | |
"""force on the muller potential | |
Parameters | |
---------- | |
position : list-like | |
x,y in a tuple or list or array | |
Returns | |
------- | |
force : np.ndarray | |
force in x direction, y direction in a length 2 numpy 1D array | |
""" | |
return np.array(F(*position)) | |
return force | |
def linear_interpolate(initial, final, n_images): | |
"""Linear interpolation between initial and final | |
Example | |
------- | |
>>> linear_interpolate([0,0], [1,1], 5) | |
[[ 0. 0. ] | |
[ 0.25 0.25] | |
[ 0.5 0.5 ] | |
[ 0.75 0.75] | |
[ 1. 1. ]] | |
""" | |
delta = np.array(final) - np.array(initial) | |
images = np.ones((n_images,) + delta.shape, dtype=np.float) | |
# goes from 0 to 1 | |
lambdas = np.arange(n_images, dtype=np.float) / (n_images-1) | |
for i in range(n_images): | |
images[i] = initial + lambdas[i]*delta | |
return images | |
def get_spring_deriv(images, spring_constant): | |
"""Calculate the spring force on each image | |
Note: the forces on the first and last bead are not calculated -- they are | |
just set to zero. | |
""" | |
# the total spring energy is | |
# E_s = sum_{i=0}^n_images (n_images * k / 2) * (R[i] - R[i-1])**2 | |
# taking the partial derivative with respect to R[j], we have | |
# n_images * k * (2*R[j] - R[j-1] - R[j+1]) | |
n_images = len(images) | |
force = np.zeros_like(images) | |
# the force on the endpoints is going to be zeroed out. | |
for j in range(1, n_images - 1): | |
force[j] = (n_images * spring_constant | |
* (2*images[j] - images[j-1] - images[j+1])) | |
return force | |
def projection(forces, images): | |
"""Given a force vector at each image, separate it out into parallel and | |
perpindicular components | |
Note: This calculation ignores the endpoints. forces[0] and forces[-1] don't | |
enter into the calculation | |
Returns | |
------- | |
parallel : np.ndarray, shape=(n_images, n_dims) | |
perpindicular : np.ndarray, shape=(n_images, n_dims) | |
""" | |
parallel, perpindicular = np.zeros_like(images), np.zeros_like(images) | |
n_images = len(images) | |
for i in range(1, n_images - 1): | |
difference = images[i+1] - images[i-1] | |
unit_tangent = difference / np.linalg.norm(difference) | |
# projection of forces[i] parallel to tangent | |
parallel[i] = unit_tangent * np.dot(forces[i], unit_tangent) | |
perpindicular[i] = forces[i] - parallel[i] | |
return Components(parallel, perpindicular) | |
def grad_descent(grad, x0, max_steps=100, tol=1e-7, step_size=0.1, callback=None): | |
x_k = x0 | |
k = 0 | |
# make the change bigger than tol so that it initializes okay | |
change = tol + 1 | |
while np.linalg.norm(change) > tol and k < max_steps: | |
if callback is not None: | |
callback(x_k, k) | |
change = step_size * grad(x_k) | |
x_k -= change | |
k += 1 | |
return x_k | |
def nudged_elastic_band(initial, final, n_intermediates, spring_constant, | |
grad, callback): | |
""" | |
Parameters | |
---------- | |
n_intermediates : int | |
number of intermediate images between the initial and final image | |
""" | |
# start with linear interpolation | |
# plus 2 for the end points | |
n_images = n_intermediates + 2 | |
images = linear_interpolate(initial, final, n_images) | |
# shape = (n_images, n_dimensions). this is the true shape | |
# size = (n_images * n_dimensions). for bfgs, we need to make stuff 1D | |
shape, size = images.shape, images.size | |
def neb_derivative(images): | |
# negative of force on each bead from the potential energy function | |
potential_deriv = grad(images) | |
# negative of force on each bead from the mutual springs | |
spring_deriv = get_spring_deriv(images, spring_constant) | |
# keep only the perpindicular comp. of the potential force and | |
# the parallel comp. of the spring force | |
neb_deriv = (projection(potential_deriv, images).perpindicular + | |
projection(spring_deriv, images).parallel) | |
print 'norm of derivative =', np.linalg.norm(neb_deriv) | |
return neb_deriv | |
#def neb_action(flattened_images): | |
# images = flattened_images.reshape(shape) | |
# return np.sum(func(images)) + ((n_images * spring_constant / 2) * | |
# np.sum((images[1:] - images[:-1])**2)) | |
print 'starting optimization' | |
xf = grad_descent(neb_derivative, images, max_steps=5000, | |
tol=1e-5, step_size=0.001, callback=callback) | |
return xf | |
def main(): | |
muller_grad = muller_grad_factory() | |
def grad(images): | |
return np.array(map(muller_grad, images)) | |
def callback(xk, k): | |
pass | |
print 'Positions', k | |
print xk | |
#if k % 250 == 0: | |
# pp.plot(xk[:,0], xk[:,1], '-o', markersize=10, label='it:%s' % k) | |
images = nudged_elastic_band([-0.563, 1.417], [0.618, 0.01], n_intermediates=4, | |
spring_constant=1, grad=grad, callback=callback) | |
plot_v() | |
pp.plot(images[:,0], images[:,1], '-ko', markersize=10, label='finish') | |
pp.legend() | |
pp.show() | |
def plot_v(minx=-1.5, maxx=1.2, miny=-0.2, maxy=2, ax=None): | |
"Plot the Muller potential" | |
grid_width = max(maxx-minx, maxy-miny) / 200.0 | |
xx, yy = np.mgrid[minx : maxx : grid_width, miny : maxy : grid_width] | |
V = muller_potential(xx, yy) | |
# clip off any values greater than 200, since they mess up | |
# the color scheme | |
if ax is None: | |
ax = pp | |
ax.contourf(xx, yy, V.clip(max=200), 40) | |
if __name__ == '__main__': | |
main() | |
print 'helloc' |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment