FTS-Ensemble.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Mar 25 17:54:10 2018

@author: ebeyerle
"""

import matplotlib
matplotlib.use('Qt5Agg')
import numpy as np
import matplotlib.pyplot as plt
import Markov_Models as mm
from scipy import stats
from scipy import interpolate
from scipy import spatial
import os
from sklearn.neighbors import KNeighborsRegressor
from sklearn.model_selection import ShuffleSplit
from sklearn.metrics import mean_squared_error

#Do a little ZTS to get a good starting configuration
path="/home/ebeyerle/Desktop/PCA/UBQN/400ns/"
os.chdir(path)
bins=50
xx,yy=np.meshgrid(np.linspace(0,180,bins),np.linspace(0,360,bins))
#mode=1
nfrs=1
nframes=50
barlist=[]
pointslist=[]
for mode in range(1,2):
    with open(path+'anly_'+str(mode)+'.dat','r') as data:
        x=[]
        y=[]
        for line in data:
            nfrs=nfrs+1
            p=line.split()
            x.append(float(p[0]))
            y.append(float(p[1]))
    print(nfrs)
    data=[np.column_stack([x,y])]
    model = mm.MSM(data)
    his0 = model.histogram(0, bins=50)
    his1 = model.histogram(1, bins=50)
    his = model.histogram(bins=50)
    ext=model.extent
    lst=np.zeros(bins*bins)
    chk=his.T/his.sum()
    counter=-1
    for i in range(len(chk[0,:])):
        for j in range(len(chk[:,0])):
            counter=counter+1
            lst[counter]=chk[i][j]
    counter=0
    fes=np.zeros((bins,bins))
    for i in range(len(chk[0,:])):
        for j in range(len(chk[:,0])):
            if (lst[counter] <= 0):
                fes[i][j]=-np.log(1/nfrs)
                counter=counter+1
            else:
                fes[i][j]=-np.log(lst[counter])
                counter=counter+1
    min1=np.min(fes)
    for i in range(bins):
        for j in range(bins):
            fes[i][j]=fes[i][j]-min1
    #plt.contourf(xx,yy,fes, 25, cmap='gnuplot')
    #plt.show()
    theta=np.linspace(0,180,bins)
    phi=np.linspace(0,360,bins)
    #Smooth the surface with k-nearest neighbors
    n_values = 100
    alpha = 2

    x0, x1 = np.meshgrid(theta,phi)
    X = np.column_stack([x0.ravel(), x1.ravel()])
    score = []
    for n_neighbor in np.arange(alpha, n_values * alpha + 1, alpha):
        for train_index, test_index in ShuffleSplit(n_splits=1, test_size=0.40, random_state=24).split(X):
            X_train, y_train = X[train_index], fes.ravel()[train_index]
            X_test, y_test = X[test_index], fes.ravel()[test_index]

        knn = KNeighborsRegressor(n_neighbors=n_neighbor).fit(X_train, y_train)
        score.append(mean_squared_error(knn.predict(X_test), y_test))
    if np.argmin(score) == 0:
        n_neighbors=2
        print("Hmm... no nearest neighbors. Setting k:=1.")
    else:
        n_neighbors = np.argmin(score) * alpha        
    print("k = {:d}".format(np.argmin(score)))
    knn = KNeighborsRegressor(n_neighbors=n_neighbors).fit(X, fes.ravel())
    fes_knn = knn.predict(X).reshape(bins, bins)
    dVy,dVx=np.gradient(fes_knn)
    dVx=interpolate.interp2d(theta,phi,dVx,kind='quintic')
    dVy=interpolate.interp2d(theta,phi,dVy,kind='quintic')
    plt.close()
    plt.contourf(xx,yy,fes,cmap='gnuplot')
    plt.quiver(xx,yy,-dVx(theta,phi),-dVy(theta,phi))
    #plt.show()
    #Now convert the code from the ZTS Matlab code
    tol1=0.00001 #Convergence criterion
    tollist=[]
    flag1=1
    nstepplot=100
    nstep=1000
    g1=np.linspace(0,1,nframes)
    #Perform the interpolation; same as above
    min1=np.where(fes == 0.0)
    ya=min1[0][0]
    xa=min1[1][0]
    xb=140
    yb=300
    x=(xb-xa)*g1+xa
    y=(x-xa)*(yb-ya)/(xb-xa)+ya
    dx=x-np.roll(x,1,axis=0)
    dy=y-np.roll(y,1,axis=0)
    dx[0]=0
    dy[0]=0
    lxy=np.cumsum(np.sqrt(dx**2+dy**2))
    lxy=lxy/lxy[nframes-1]
    #Honestly not sure why these two lines are here; the x, y that are returned
    #are identical to the x, y generated by the linear interpolation.
    x=interpolate.interp1d(lxy,x)(g1)
    y=interpolate.interp1d(lxy,y)(g1)
    xi=x
    yi=y
    #plt.plot(xi,yi,'w',lw=1)
    plt.scatter(xi,yi,marker='x',c='w')
    plt.plot(x,y,'b')
    counter=1
    #plt.savefig("/home/ebeyerle/Desktop/fifty_fes_animate_mode"+str(mode)+"frame00"+str(counter)+".jpeg")
    plt.close()
    #plt.show()
    laplist=[]
    grady,gradx=np.gradient(fes_knn)
    lapxy,lapxx=np.gradient(gradx)
    lapyy,lapyx=np.gradient(grady)
    h=1/(np.array([lapxx.max(),lapxy.max(),lapyx.max(),lapyy.max()]).max())
    #h=0.1#Just a parameter -- proabably need to adjust to get good convergence
    print('h=',h)
    for i in range(nstep):
        dVy,dVx=np.gradient(fes_knn)
        dVx=interpolate.interp2d(theta,phi,dVx,kind='quintic')
        dVy=interpolate.interp2d(theta,phi,dVy,kind='quintic')
        x0=x
        y0=y
        # Evolve the string using steepest descent
        for j in range(len(x)):
            x[j]=x[j]-h*dVx(x[j],y[j])[0]
            y[j]=y[j]-h*dVy(x[j],y[j])[0]
            #Enforce PBCs
            #KLUDGE: Not letting them cross for now
            if y[j] <=0.0: 
                y[j]=0.0
            if y[j] >=360.0:
                y[j]=360.0
            if x[j] <=0.0:
                x[j] = 0.0
            if x[j] >= 180.0:
                x[j]=180.0
        #Reparameterize
        dx=x-np.roll(x,1,axis=0)
        dy=y-np.roll(y,1,axis=0)
        dx[0]=0
        dy[0]=0
        #Find the string length...
        lxy=np.cumsum(np.sqrt(dx**2+dy**2))
        lxy=lxy/lxy[nframes-1]
        #...and place each state at an equally-spaced length along the string
        x=interpolate.interp1d(lxy,x)(g1)
        y=interpolate.interp1d(lxy,y)(g1)
        if flag1 == 1 and i%nstepplot == 0:
            counter=counter+1
            #plt.plot(x,y)
            if counter < 10:
                plt.contourf(xx,yy,fes,cmap='gnuplot')
                plt.quiver(xx,yy,-dVx(theta,phi),-dVy(theta,phi))
                plt.scatter(x,y,marker='o',c='w')
                plt.plot(x,y,'b')
                #plt.savefig("/home/ebeyerle/Desktop/fifty_fes_animate_mode"+str(mode)+"frame00"+str(counter)+".jpeg")
                plt.close()
            elif counter >= 10 and counter <100:
                plt.contourf(xx,yy,fes,cmap='gnuplot')
                plt.quiver(xx,yy,-dVx(theta,phi),-dVy(theta,phi))
                plt.scatter(x,y,marker='o',c='w')
                plt.plot(x,y,'b')
                #plt.savefig("/home/ebeyerle/Desktop/fifty_fes_animate_mode"+str(mode)+"frame0"+str(counter)+".jpeg")
                plt.close()
            else:
                plt.contourf(xx,yy,fes,cmap='gnuplot')
                plt.quiver(xx,yy,-dVx(theta,phi),-dVy(theta,phi))
                plt.scatter(x,y,marker='o',c='w')
                plt.plot(x,y,'b')
                #plt.savefig("/home/ebeyerle/Desktop/fifty_fes_animate_mode"+str(mode)+"frame"+str(counter)+".jpeg")
                plt.close()
        #    plt.show()
        tol=(np.linalg.norm(x-x0)+np.linalg.norm(y-y0))/(nframes)
        tollist.append(tol)
        #if len(tollist) >=3:
        #    if tollist[i-3]-tollist[i] == 0:
        #        print('converged')
        #        break
        if tol <= tol1: #Has the string converged?
            break #It has converged!
        if i%nstepplot == 0:
            print(tol)
    print('\n')
    print('\n')
    print('ZTS calculation with ',nframes,' images.\n')
    if tol >= tol1:
        print('The method did not converge in ',nstep,' steps.')
    else:
        print('The method converged after ',nstep,' steps.')
    #plt.plot(x,y)
    plt.contourf(xx,yy,fes,cmap='gnuplot')
    plt.quiver(xx,yy,-dVx(theta,phi),-dVy(theta,phi))
    plt.scatter(x,y,marker='o',c='w')
    plt.plot(x,y,'b')
    plt.show()


    bins=50
    xx,yy=np.meshgrid(np.linspace(0,180,bins),np.linspace(0,360,bins))
    #mode=1
    nfrs=1
    nframes=50
    kT=0.5 #Unitless temperature
    barlist=[]
    pointslist=[]
    with open(path+'anly_'+str(mode)+'.dat','r') as data:
        x=[]
        y=[]
        for line in data:
            nfrs=nfrs+1
            p=line.split()
            x.append(float(p[0]))
            y.append(float(p[1]))
    print(nfrs)
    data=[np.column_stack([x,y])]
    model = mm.MSM(data)
    his0 = model.histogram(0, bins=50)
    his1 = model.histogram(1, bins=50)
    theta=np.linspace(0,180,bins)
    phi=np.linspace(0,360,bins)
    #fig=plt.figure(figsize=(12,8))
    dVy,dVx=np.gradient(fes)
    dVx=interpolate.interp2d(theta,phi,dVx,kind='quintic')
    dVy=interpolate.interp2d(theta,phi,dVy,kind='quintic')
    plt.close()
    plt.contourf(xx,yy,fes,cmap='gnuplot')
    plt.xlabel(r'$\theta (deg)$')
    plt.ylabel(r'$\phi (deg)$')
    plt.quiver(xx,yy,-dVx(theta,phi),-dVy(theta,phi))
    #plt.show()
    #Now convert the code from the ZTS Matlab code
    tol1=0.00001 #Convergence criterion
    tollist=[]
    flag1=1
    nstepplot=1000
    nstep=10000
    g1=np.linspace(0,1,nframes)
    #Perform the interpolation; same as above
    ya=y[0]
    xa=x[0]
    xb=x[-1]
    yb=y[-1]
    x=(xb-xa)*g1+xa
    y=(x-xa)*(yb-ya)/(xb-xa)+ya
    dx=x-np.roll(x,1,axis=0)
    dy=y-np.roll(y,1,axis=0)
    dx[0]=0
    dy[0]=0
    lxy=np.cumsum(np.sqrt(dx**2+dy**2))
    lxy=lxy/lxy[nframes-1]
    #Honestly not sure why these two lines are here; the x, y that are returned
    #are identical to the x, y generated by the linear interpolation.
    x=interpolate.interp1d(lxy,x)(g1)
    y=interpolate.interp1d(lxy,y)(g1)
    xi=x
    yi=y
    #plt.plot(xi,yi,'w',lw=1)
    plt.scatter(xi,yi,marker='x',c='w')
    plt.plot(x,y,'b')
    counter=1
    #plt.savefig("/home/ebeyerle/Desktop/fifty_fes_animate_mode"+str(mode)+"frame00"+str(counter)+".jpeg")
    plt.close()
    #plt.show()
    laplist=[]
    grady,gradx=np.gradient(fes)
    lapxy,lapxx=np.gradient(gradx)
    lapyy,lapyx=np.gradient(grady)
    h=1/(np.array([lapxx.max(),lapxy.max(),lapyx.max(),lapyy.max()]).max())
    n_strings=10
    #h=0.1#Just a parameter -- proabably need to adjust to get good convergence
    print('h=',h)
    for k in range(nstep):
        etax=np.random.randn(1,nframes)
        etay=np.random.randn(1,nframes)
        dVy,dVx=np.gradient(fes)
        dVx=interpolate.interp2d(theta,phi,dVx,kind='quintic')
        dVy=interpolate.interp2d(theta,phi,dVy,kind='quintic')
        x0=x
        y0=y
        # Evolve the string using steepest descent
        xstrings=x+np.random.normal(scale=np.sqrt(2*h*kT),size=(n_strings,nframes))
        ystrings=y+np.random.normal(scale=np.sqrt(2*h*kT),size=(n_strings,nframes))
        #avgx=[]
        #avgy=[]
        sumx=np.zeros(len(x))
        sumy=np.zeros(len(y))
        for j in range(len(x)):
            for i in range(n_strings):
                xstrings[i][j]=xstrings[i][j]-h*dVx(xstrings[i][j],ystrings[i][j])[0]+np.sqrt(2*h*kT)*etax[0][j]
                ystrings[i][j]=ystrings[i][j]-h*dVy(xstrings[i][j],ystrings[i][j])[0]+np.sqrt(2*h*kT)*etay[0][j]
                #Enforce PBCs
                #KLUDGE: Not letting them cross for now
                if ystrings[i][j] <=0.0: 
                    ystrings[i][j]=0.0
                if ystrings[i][j] >=360.0:
                    ystrings[i][j]=360.0
                if xstrings[i][j] <=0.0:
                    xstrings[i][j] = 0.0
                if xstrings[i][j] >= 180.0:
                    xstrings[i][j]=180.0
                sumx[j]=sumx[j]+xstrings[i][j]
                sumy[j]=sumy[j]+ystrings[i][j]
            sumx[j]=sumx[j]/n_strings
            sumy[j]=sumy[j]/n_strings
            #avgx.append(sumx[j])
            #avgy.append(sumy[j])
        x=sumx#np.array(avgx)
        y=sumy#np.array(avgy)
        #Reparameterize
        dx=x-np.roll(x,1,axis=0)
        dy=y-np.roll(y,1,axis=0)
        dx[0]=0
        dy[0]=0
        #Find the string length...
        lxy=np.cumsum(np.sqrt(dx**2+dy**2))
        lxy=lxy/lxy[nframes-1]
        #...and place each state at an equally-spaced length along the string
        x=interpolate.interp1d(lxy,x)(g1)
        y=interpolate.interp1d(lxy,y)(g1)
        #if flag1 == 1 and i%nstepplot == 0:
        #    counter=counter+1
            #plt.plot(x,y)
        #    if counter < 10:
        #        plt.contourf(xx,yy,fes,cmap='gnuplot')
        #        plt.xlabel(r'$\theta (deg)$')
        #        plt.ylabel(r'$\phi (deg)$')
        #        plt.quiver(xx,yy,-dVx(theta,phi),-dVy(theta,phi))
        #        plt.scatter(x,y,marker='o',c='w')
        #        plt.plot(x,y,'b')
                #plt.savefig("/home/ebeyerle/Desktop/FTS_animate_mode"+str(mode)+"frame00"+str(counter)+".jpeg")
        #        plt.close()
        #    elif counter >= 10 and counter <100:
        #        plt.contourf(xx,yy,fes,cmap='gnuplot')
        #        plt.xlabel(r'$\theta (deg)$')
        #        plt.ylabel(r'$\phi (deg)$')
        #        plt.quiver(xx,yy,-dVx(theta,phi),-dVy(theta,phi))
        #        plt.scatter(x,y,marker='o',c='w')
        #        plt.plot(x,y,'b')
                #plt.savefig("/home/ebeyerle/Desktop/FTS_animate_mode"+str(mode)+"frame0"+str(counter)+".jpeg")
        #        plt.close()
        #    else:
        #        plt.contourf(xx,yy,fes,cmap='gnuplot')
        #        plt.xlabel(r'$\theta (deg)$')
        #        plt.ylabel(r'$\phi (deg)$')
        #        plt.quiver(xx,yy,-dVx(theta,phi),-dVy(theta,phi))
        #        plt.scatter(x,y,marker='o',c='w')
        #        plt.plot(x,y,'b')
                #plt.savefig("/home/ebeyerle/Desktop/FTS_animate_mode"+str(mode)+"frame"+str(counter)+".jpeg")
        #        plt.close()
        #    plt.show()
        tol=(np.linalg.norm(x-x0)+np.linalg.norm(y-y0))/(nframes)
        tollist.append(tol)
        #if len(tollist) >=3:
        #    if tollist[i-3]-tollist[i] == 0:
        #        print('converged')
        #        break
        if tol <= tol1: #Has the string converged?
            break #It has converged!
        if k%nstepplot == 0:
            print(tol)
    print('\n')
    print('\n')
    print('FTS calculation with ',nframes,' images.\n')
    if tol >= tol1:
        print('The method did not converge in ',nstep,' steps.')
    else:
        print('The method converged after ',nstep,' steps.')
    #plt.plot(x,y)
    plt.contourf(xx,yy,fes,cmap='gnuplot')
    plt.xlabel(r'$\theta (deg)$')
    plt.ylabel(r'$\phi (deg)$')
    plt.quiver(xx,yy,-dVx(theta,phi),-dVy(theta,phi))
    plt.scatter(x,y,marker='o',c='w')
    plt.plot(x,y,'b')
    plt.show()
    #Generate tangent line
    tx=np.roll(x,-1,axis=0)-np.roll(x,1,axis=0)
    ty=np.roll(y,-1,axis=0)-np.roll(y,1,axis=0)
    #Use trapezoidal integration
    from scipy import integrate
    Vz=integrate.cumtrapz(tx*dVx(x,y)+ty*dVy(x,y))
    Vz=0.5*Vz
    Vz=Vz-np.min(Vz)
    ntxy=np.sqrt(tx*tx+ty*ty)
    tx=tx/ntxy
    ty=ty/ntxy
    err=np.trapz(1-(tx*dVx(x,y)+ty*dVy(x,y))**2/(dVx(x,y)*dVx(x,y)+dVy(x,y)*dVy(x,y))/nframes)
    print('\n')
    #print('Estimate of the difference between the discretized MEP and the actual MEP: ',err,'.')
    #Plot the 1D FES along the string
    ifes=interpolate.interp2d(theta,phi,fes,kind='quintic')
    points=[]
    for i in range(len(x)):
        points.append(ifes(x[i],y[i])[0])
        #print(points[i])
    pointslist.append(points)
    np.savetxt("points"+str(mode),points)
    np.savetxt("coords_"+str(mode)+".dat",np.column_stack([x,y]))
    plt.plot(points,'k')
    plt.xlabel('Image Number')
    plt.ylabel('free energy ($k_BT$)')
    bar=np.max(np.array(points))-np.min(np.array(points))
    print('Barrier Height ($k_BT$): ',bar)
    barlist.append(bar)
    np.savetxt("barlist.dat",np.array(barlist))
    np.savetxt("coords_"+str(mode)+".dat",np.column_stack([x,y]))