2D PMF from NAMD

plot2d.py

import sys
import mdtraj as md
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
from matplotlib.mlab import griddata
import matplotlib.pyplot as plt
def Get1DPES(df):
    XD = []
    YD = []
    ZD = []
    groups = df.groupby('X').groups
    for g in np.sort(list(groups.keys())):
        g_df = df[df.X==g]
        XD.append(g)
        YD.append(g_df[g_df.Z == g_df.Z.min()].Y.values[0])
        ZD.append(g_df.Z.min())
    return (XD,YD,ZD)
def AdvCombo2DPlot(df):
    XD,YD,ZD=Get1DPES(df)    
    dy = 0.1
    dx = 0.1
    y, x = np.mgrid[slice(df.Y.min(), df.Y.max() + dy, dy),
                    slice(df.X.min(), df.X.max() + dx, dx)]

    z = griddata(df.X, df.Y, df.Z, x, y, interp='linear')
#    z = griddata(df.X, df.Y, df.Z, x, y, interp='linear')

    # x and y are bounds, so z should be the value *inside* those bounds.
    # Therefore, remove the last value from the z array.



    # pick the desired colormap, sensible levels, and define a normalization
    # instance which takes data values and translates those into levels.
    cmap = plt.get_cmap('bwr')

    fig, (ax0, ax1) = plt.subplots(nrows=2,figsize=(10,10),sharex=True)

    im = ax1.scatter(XD, ZD, c=YD, cmap=plt.cm.gnuplot)
    cbbar = fig.colorbar(im, ax=ax1)
    cbbar.set_label('RMSD ($\AA$)')
    ax1.set_ylabel('Energy (kcal/mol)')
    ax1.set_xlabel('Drug-BSite COM distance ($\AA$)')
#    ax1.set_title('UN')


    # contours are *point* based plots, so convert our bound into point
    # centers
    cf=ax0.contourf(x,y, z, cmap=cmap)
    CS = ax0.contour(x, y, z, 6,
                 colors='k',  # negative contours will be dashed by default
                 )
    plt.clabel(CS, fontsize=9, inline=1)
    cabar = fig.colorbar(cf, ax=ax0)
    cabar.set_label('Energy (kcal/mol)')
    ax0.plot(XD,YD,'k.-')
    ax0.set_ylabel('RMSD ($\AA$)')
    fig.tight_layout()
    plt.savefig('Comp2D.png',dpi=300)
    return None
df = pd.read_csv(sys.argv[1],delim_whitespace=True,skip_blank_lines=True,comment='#',header=None)
df.columns=['X','Y','Z']
df=df[df.Z!=df.Z.max()]
#print(Get1DPES(df))
AdvCombo2DPlot(df)