Save2traj

Description

WARNING: This script is still under development so please use at your own risk!

This script can be used to generate a DCD trajectory file after you have loaded multiple states into a single object.

Currently, the only supported trajectory format is DCD but there will be support for other formats in the future.

USAGE

load the script using the run command

save2traj (selection, name)

"name" corresponds to the output name and ".dcd" will be appended to the end of the name automatically.

EXAMPLES

save2traj 1bna & c. A+B, output

This will select chains A and B from the object 1bna and save the selection from each state into the trajectory file called "output.dcd".

Script

WARNING: This script is still under development so please use at your own risk!

from pymol import cmd
from pymol import stored
import struct

def save2traj (selection,name,format="dcd"):

    #Author: Sean Law
    #Michigan State University

    #Get NATOMS, NSTATES
    NATOMS=cmd.count_atoms(selection)
    NSTATES=cmd.count_states()

    #Determine Trajectory Format
    format=format.lower()
    if (format == "charmm"):
        name=name+".dcd"
    elif (format == "amber"):
        print "The amber format has not been implemented yet"
        return
        name=name+".trj"
    elif (format == "trj"):
        print "The amber format has not been implemented yet"
        return
        name=name+".trj"
        format="amber"
    else:
        name=name+".dcd"
        format="charmm"

    f=open(name,'wb')

    #Write Trajectory Header Information
    if (format == "charmm"):
        writeCHARMMheader(f,NSTATES,NATOMS)
    elif (format == "amber"):
        print "The amber format has not been implemented yet"
        return
    else:
        print "Unknown format"
        return

    #Write Trajectory Coordinates
    if (format == "charmm"):
        fmt=str(NATOMS)+'f'
        for state in range(cmd.count_states()):
            stored.xyz=[]
            cmd.iterate_state(state,selection,"stored.xyz.append([x,y,z])")
            for xyz in range (3):
                coor=[]
                for atom in range (NATOMS):
                    coor.append(stored.xyz[atom][xyz])
                writeFortran(f,coor,fmt,length=NATOMS*4)
    elif (format == "amber"):
        print "The amber format has not been implemented yet"
        return
    else:
        print "Unknown format"
        return

    f.close()
    return
cmd.extend("save2traj",save2traj)

def writeCHARMMheader (f,NSTATES,NATOMS):
    header=['CORD']
    fmt='4s9i1f10i'
    icontrol=[]
    for i in range(20):
        #Initialize icontrol
        icontrol.insert(i,0)
    icontrol[0]=NSTATES
    icontrol[1]=1 
    icontrol[2]=1
    icontrol[3]=NSTATES
    icontrol[7]=NATOMS*3-6
    icontrol[9]=2.045473
    icontrol[19]=27

    for i in range(20):
        header.append(icontrol[i])

    writeFortran(f,header,fmt)

    #Title
    fmt='i80s80s'
    title=[2]
    title.append('* TITLE')
    while (len(title[1])<80):
        title[1]=title[1]+' '
    title.append('* Generated by savetraj.py (Author: Sean Law)')
    while (len(title[2])< 80):
        title[2]=title[2]+' '
    
    writeFortran(f,title,fmt,length=160+4)

    #NATOM
    fmt='i'
    writeFortran(f,[NATOMS],fmt)

    return
cmd.extend("writeCHARMMheader",writeCHARMMheader)

def readtraj (name):
    f=open(name,'rb')
    #Read Header
    fmt='4s9i1f10i'
    header=readFortran(f,fmt)

    #Read Title
    readCHARMMtitle(f)

    #Read NATOMS
    fmt='i'    
    [NATOMS]=readFortran(f,fmt)
    
    #Read COORDINATES
    

    f.close()
    return
cmd.extend("readtraj",readtraj)

def writeFortran (f,buffer,fmt,length=0):
    
    if (length == 0):
        length=len(buffer)*4
    f.write(struct.pack('i',length)) #Fortran unformatted
    f.write(struct.pack(fmt,*(buffer)))
    f.write(struct.pack('i',length)) #Fortran unformatted

    return
cmd.extend("writeFortran",writeFortran)

def readFortran (f,fmt):
    [bytes]=struct.unpack('i',f.read(4))
    buffer=struct.unpack(fmt,f.read(bytes))
    [bytes]=struct.unpack('i',f.read(4))

    return buffer
cmd.extend("readFortran",readFortran)

def readCHARMMtitle(f):
    [bytes]=struct.unpack('i',f.read(4))
    [lines]=struct.unpack('i',f.read(4))
    fmt=''
    for line in range(lines):
        fmt=fmt+'80s'
    buffer=(struct.unpack(fmt,f.read(80*lines)))
    [bytes]=struct.unpack('i',f.read(4))
    #print buffer
    return
cmd.extend("readCHARMMtitle",readCHARMMtitle)