Nudged Elastic Band in Python

neb.py

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'