Library of Protein analysis module for me!!

lProam.py

# GPL v 3.0 
# python library name: library of Protein analysis module(lProam)
# lProam.py
# written by Naoyuki Miyashita on 12/18/2013 
# Naoyuki Miyashita (c) 2013
#
#-----------------------------------
# please install 
# prody, pylab(matplotlib), numpy and scipy,
# before you use this scripts.
#-----------------------------------
# ussage:
#-----------------------------------
# Please use prody selection rules in the atom selections.
#.
from prody import *
from pylab import *
import numpy as np
import scipy as sp
import math
#

def Constants():
    kboltz = 1.987191*10**(-3)
    jkboltz = 1.380662*10**(-23)
    avogadro= 6.022045*10**(23)
    jkcal = 4186.05
    kcalmol = jkcal/avogadro
    return(kboltz, jkboltz,avogadro,jkcal,kcalmol)

def file_read(filename):
    pdb1=parsePDB(filename)
    return(pdb1)

def dcd_read(filename):
    dcd = parseDCD(filename)
    return(dcd)

def dcd_trj(filename):
    trj = Trajectory(filename)
    return(trj)

def readFrame(filepdb,filedcd):
    pdb        = file_read(filepdb)
    dcd        = dcd_read(filedcd)
    dcd.setAtoms(pdb)
    dcd.setCoords(pdb)
    return (pdb,dcd)

def readfile(filename):
    f=open(filename,"r")
    box=f.readlines()
    f.close()
    return(box)

def devide2Line(f):
    lines=f.replace("\n","")
    lines=lines.strip()
    return (lines)

def splitData(line):
    lines=line.split()
    return (lines)

def mkArray2(box, col1,col2):
    nline=len(box)
    jobids=([])
    jobnames=([])
    for i in range(nline):
        line=box[i]
        line2=devide2Line(line)
        if len(line2) == 0:
           continue
        lines=splitData(line2)
        jobids.append( lines[col1] )
        jobnames.append( lines[col2] )
    job_array=(jobids,jobnames)
    return (job_array)

def ckKeyWords(job_array, keyword):
    (jobids,jobnames) = job_array
    FindJobID = "non"
    for i in range(len(jobids)):
        if jobnames[i] == keyword:
           FindJobID = jobids[i]
           break
    return (FindJobID)

def getSelectedFrames(pdb,dcd,key_select):
    key_pdb    = pdb.select(key_select)
    dcd.setAtoms(key_pdb)
    return (pdb,dcd)

def moveFrame2Center(pdb,dcd):
    #set_center(pdb,key_select,key)
    a=calcCenter(pdb).round(3)
    print a
    a=np.array([0,0,0])
    moveAtoms(pdb,to=zeros(3))
    #dcd.setAtoms(pdb)
    #dcd.setCoords(pdb)
    for i, frame in enumerate(dcd):
        a= calcCenter(frame).round(3)
        #print a
        coor = frame.getCoords()
        pdb_coor = copy(coor)
        print coor
        print "chk"
        #for j, r in enumerate(coor):
        for j in range(len(coor)):
            r = transSel(coor[j],a,'xyz')
            pdb_coor[j] = r
            #print 'r=',r
        print r
        frame.setCoords(pdb_coor)
    #    moveAtoms(frame,to=zeros(3))
    ##dcd.superpose()
    return (pdb,dcd)

def transSel(coord,center,key):
    if key == 'xyz':
       coord[0] = coord[0]-center[0]
       coord[1] = coord[1]-center[1]
       coord[2] = coord[2]-center[2]
    elif key == 'xy':
       coord[0] = coord[0]-center[0]
       coord[1] = coord[1]-center[1]
    elif key == 'xz':
       coord[0] = coord[0]-center[0]
       coord[2] = coord[2]-center[2]
    elif key == 'yz':
       coord[1] = coord[1]-center[1]
       coord[2] = coord[2]-center[2]
    elif key == 'x':
       coord[0] = coord[0]-center[0]
    elif key == 'y':
       coord[1] = coord[1]-center[1]
    elif key == 'z':
       coord[2] = coord[2]-center[2]
    return (coord)

def install_frames(filepdb,filedcd,key_select):
    pdb        = file_read(filepdb)
    dcd        = dcd_read(filedcd)
    dcd.setAtoms(pdb)
    dcd.setCoords(pdb)
    key        = pdb.select(key_select)
    dcd.setAtoms(key)
    dcd.superpose()
    return (dcd)


def file_write(filename,pdb):
    writePDB(filename,pdb)
    return()

def set_chain(pdb,chid):
    pdb_ch  = pdb.getChids()
    pdb_ch_tmp=[]
    for i in range(len(pdb_ch)):
        pdb_ch[i]=chid
    pdb.setChids(pdb_ch)
    return(pdb) 

def set_trans(pdb,center,key):
    pdb_coords = pdb.getCoords()
    for i in range(len(pdb_coords)):
        dist2 = pdb_coords[i]
        if key == 'x':
           dist2[0]=dist2[0]-center[0]
        elif key == 'y':
           dist2[1]=dist2[1]-center[1]
        elif key == 'z':
           dist2[2]=dist2[2]-center[2]
        elif key == 'xy':
           dist2[0]=dist2[0]-center[0]
           dist2[1]=dist2[1]-center[1]
        elif key == 'xz':
           dist2[0]=dist2[0]-center[0]
           dist2[2]=dist2[2]-center[2]
        elif key == 'yz':
           dist2[1]=dist2[1]-center[1]
           dist2[2]=dist2[2]-center[2]
        else:
           dist2[0]=dist2[0]-center[0]
           dist2[1]=dist2[1]-center[1]
           dist2[2]=dist2[2]-center[2]
        pdb_coords[i] =  dist2 
    pdb.setCoords(pdb_coords)
    return(pdb)

def getCoords(pdb,selec):
    key_sel    = pdb.select(selec)
    coords     = key_sel.getCoords()
    return(coords)

def set_center(pdb,center_selec,key):
    center = get_center(pdb,center_selec)
    set_trans(pdb,center,key)
    return(pdb)

def get_center(pdb,center_selec):
    selec_c   = pdb.select(center_selec)
    center    = calcCenter(selec_c)
    return(center)

def get_pdb_center(filename,center_selec):
    #read pdb
    center = np.array([])
    pdb1   = file_read(filename)
    key_pdb    = pdb1.select(center_selec)
    #setup pdb
    cent = center_selec + " and name CA"
    center = get_center(key_pdb,center_selec)
    # copy pdb
    return(center)

def setup_pdb(filename,center_selec,key):
    #read pdb
    pdb1   = file_read(filename)
    key_pdb    = pdb1.select(center_selec)
    #setup pdb
    cent = center_selec + " and name CA"
    pdb1   = set_center(key_pdb,center_selec,'xy')
    # copy pdb
    pdb1   = set_chain(pdb1,key)
    return(pdb1)


def rotation_matrix(ax,t):
    ax    = normVec(ax)
    a     = np.cos(t/2)
    b,c,d = -ax*np.sin(t/2)
    mat   = [[a*a+b*b-c*c-d*d, 2*(b*c-a*d), 2*(b*d+a*c)],
             [2*(b*c+a*d),   a*a+c*c-b*b-d*d, 2*(c*d-a*b)],
             [2*(b*d-a*c),  2*(c*d+a*b), a*a+d*d-b*b-c*c]]
    return(mat)

def normVec(ax):
    return ax/np.sqrt(np.dot(ax,ax))

def setRotAxis(pdb,angle,key):
    if key=='x':
       ax=np.array([1,0,0])
    elif key == 'y':
       ax=np.array([0,1,0])
    elif key == 'z':
       ax=np.array([0,0,1])
    elif key=='-x':
       ax=np.array([-1,0,0])
    elif key == '-y':
       ax=np.array([0,-1,0])
    elif key == '-z':
       ax=np.array([0,0,-1])
    else:
       print "NOT Work on getRotAxis, use getRot\n"
    pdb=setRot(pdb,angle,ax)
    return(pdb)
    
def setRot(pdb,angle,ax):
    rmat=rotation_matrix(ax,angle)
    pdb_coords=pdb.getCoords()
    trmat = np.transpose(rmat)
    pdb_coords=np.dot(pdb_coords, trmat)
    pdb.setCoords(pdb_coords)
    return(pdb)

def getRadii(coord):
    return np.sqrt(np.dot(coord,coord))

def getTheta(ax,order,r):
    d=ax/r
    if order>0:
        t=arccos(d)
    else:
        t=-arccos(d)
    return t

def getAngleX(coord1,coord2,origin_vec):
    coord1 = coord1/np.linalg.norm(coord1)
    coord2 = coord2/np.linalg.norm(coord2)
    origin_vec = origin_vec/np.linalg.norm(origin_vec)
    dot = np.dot(coord1,coord2)
    vec1to2 = np.cross(coord1,coord2)
    vec1to2 = vec1to2/np.linalg.norm(vec1to2)
    order = np.dot(vec1to2, origin_vec)
    angle = getTheta(dot,order,1)
    return(angle) 

def getAngle(coord,key):
    if key=="xy":
        coord[2]=0.0
        ax      = coord[0]
        order   = coord[1]
    elif key=="yz":
        coord[0]=0.0
        ax      = coord[1]
        order   = coord[2]
    elif key=="xz":
        coord[1]=0.0
        ax      = coord[2]
        order   = coord[0]
    else:
        ax      = coord[2]
        order   = coord[2]
    r=getRadii(coord)
    t=getTheta(ax,order,r)
    return t

def setTilt0(pdb,base_selec, sub_selec):
    center     = get_center(pdb,base_selec)
    #
    pdb        = set_center(pdb,base_selec,'xyz')
    point_sub  = get_center(pdb,sub_selec)
    point_base = get_center(pdb,base_selec)
    tmp_sub    = np.copy(point_sub)
    t          = getAngle(tmp_sub,'xy')
    pdb        = setRotAxis(pdb,t,'z')
    point_sub2 = get_center(pdb,sub_selec)
    tmp_sub    = np.copy(point_sub2)
    g          = getAngle(tmp_sub,'xz')
    # chk up down
    val        = point_sub[2] - point_base[2]
    if val < 0:
        gpi=g + np.pi
    else:
        gpi=g
    pdb        = setRotAxis(pdb,-gpi,'-y')
    point_sub3 = get_center(pdb,sub_selec)
    point_c2   = get_center(pdb,base_selec)
    pdb        = set_trans(pdb,-center,'z')
    #
    return(pdb)

def setRho0(pdb,base_selec,key_selec,addPi):
    center     = get_center(pdb,base_selec)
    #
    pdb        = set_center(pdb,base_selec,'xyz') # base center
    point_key  = get_center(pdb,key_selec)
    point_base = get_center(pdb,base_selec) # check 0,0,0
    tmp_key    = np.copy(point_key)
    t          = getAngle(tmp_key,'xy')
    #t          = t+np.pi
    t          = t+ addPi
    d_t        = np.degrees(t)
    print 'Initial Angle rho1=',d_t
    pdb        = setRotAxis(pdb,t,'z')
    pdb        = set_trans(pdb,-center,'z') # revert
    return(pdb)

def set0(pdb,base_selec,sub_selec,key_selec):
    pdb = setTilt0(pdb,base_selec, sub_selec)
    pdb = setRho0(pdb,base_selec,key_selec)
    return(pdb)

def setup_rotate_xang_rho(pdb,alpha,rho):
    # dummy1 rho rotation
    r_rho1 =np.radians(rho)
    setRotAxis(pdb,r_rho1,'z')
    # dummy1 angle rotation
    r_alpha=np.radians(alpha)
    setRotAxis(pdb,r_alpha/2,'x')
    return(pdb)

def getMomentTensolVec(dcd):
    #
    #   a11   a12   a13
    #   a21   a22   a23
    #   a31   a32   a33
    # a11 = sum m (y^2 + z^2), a12 = -sum m yx        , a13 = -sum m zx
    # a21 = -sum m xy        , a22 = sum m (x^2 + z^2), a23 = -sum m zy
    # a31 = -sum m xz        , a32 = -sum m yz        , a33 = sum m (x^2 + y^2)
    #
    a_trj = np.array([])
    a_eig = np.array([])
    for i, frame in enumerate(dcd):
        #a = np.array([[0,0,0],[0,0,0],[0,0,0]])
        a = [[0,0,0],[0,0,0],[0,0,0]]
        a = sp.mat(a)
        center = calcCenter(frame)
        #print center
        coor = frame.getCoords()
        for j, r in enumerate(coor):
            x = r[0]-center[0]
            y = r[1]-center[1]
            z = r[2]-center[2]
            a[0,0] = a[0,0] + y**2 + z**2
            a[0,1] = a[0,1] + y*x
            a[0,2] = a[0,2] + z*x
            a[1,0] = a[1,0] + x*y
            a[1,1] = a[1,1] + x**2 + z**2
            a[1,2] = a[1,2] + z*y
            a[2,0] = a[2,0] + x*z
            a[2,1] = a[2,1] + y*z
            a[2,2] = a[2,2] + x**2 + y**2
        (eig,vec)=sp.linalg.eig(a)
        #(eig,vec)=np.linalg.eig(a)
        a_eig  = np.append(a_eig,eig)
        a_trj  = np.append(a_trj,vec)
    tot=len(a_trj)
    vtot = tot / 3
    nv   = vtot / 3
    t = a_trj.reshape((nv,3,3))
    e = a_eig.reshape((nv,3))
    return(e,t)

def getRMSD(dcd):
    rmsd = dcd.getRMSDs()
    return (rmsd)

def getRGYR(dcd):
    rgyr = zeros(len(dcd))
    for i, frame in enumerate(dcd):
        rgyr[i] = calcGyradius(frame)
    return(rgyr)


# matplotlib
def setBins(ddata,min_time,max_time):
    d = len(ddata)
    dbin = (max_time - min_time) /float(d)
    real_min_time = min_time + dbin
    real_max_time = max_time + dbin
    bins = np.arange(real_min_time,real_max_time,dbin)
    return(bins)

def showPlot(bins,out_data,xlabel,ylabel):
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plot(bins,out_data, color="black",marker="o")
    show()
    return()

def dist1D(param_set):
    (aa,x_min,x_max,nx_bin,swch)=param_set
    if swch == "hist":
        dens=False
    else:
        dens=True
    #1D histogram
    hist,bins = np.histogram(aa, bins=np.linspace(x_min,x_max,num=nx_bin), density=dens)
    return(hist,bins)

def showPlot1(param_set):
    (hist,bins,x_lab,y_lab)=param_set
    #1D histogram
    plt.xlabel(x_lab)
    plt.ylabel(y_lab)
    #
    plt.plot(bins[:-1],hist, color="black", marker="o")
    return()

def outputPNG(param_set):
    (out_prefix)=param_set
    out_fig=out_prefix + ".png"
    plt.savefig(out_fig,format = 'png', dpi=300)
    plt.show()
    plt.close()
    return()