save2traj.py

#!/usr/bin/env python
# -*- coding: UTF8 -*-

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)
    frames=header[1]

    #Read Title
    readCHARMMtitle(f)

    #Read NATOMS
    fmt='i'    
    [NATOMS]=readFortran(f,fmt)
    
    #Read COORDINATES
    for frame in range(frames):
        fmt=str(NATOMS)+'f'
        x=readFortran(f,fmt)
        y=readFortran(f,fmt)
        z=readFortran(f,fmt)

    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)

trajselect.py

#!/usr/bin/env python
# -*- coding: UTF8 -*-

# Author: Guillaume Bouvier -- guillaume.bouvier@pasteur.fr
# https://research.pasteur.fr/en/member/guillaume-bouvier/
# 2018-11-19 16:26:10 (UTC+0100)

import __main__
__main__.pymol_argv = [ 'pymol', '-cqi']
import pymol
import pymol.cmd as cmd
pymol.finish_launching()
from mdtraj.formats import DCDTrajectoryFile
import os
import sys

try:
    PDBFILENAME = sys.argv[1]
    TRJFILENAME = sys.argv[2]
    SELECTION = sys.argv[3]
    OUTBASEFILENAME = sys.argv[4]
    if len(sys.argv) > 5:
        START = int(sys.argv[5])
        STOP = int(sys.argv[6])
    else:
        START = 1
        STOP = -1
except IndexError:
    print "Help:"
    print "    pdbselect input.pdb input_traj_file 'selection string' output_base_filename [start] [stop]"
    sys.exit()

CHUNK_SIZE = 100

with DCDTrajectoryFile('%s.dcd'%OUTBASEFILENAME, 'w') as f:
    while True:
        STOPCHUNK = START + CHUNK_SIZE - 1
        cmd.load(PDBFILENAME, object='traj')
        cmd.load_traj(TRJFILENAME, object='traj', start=START, stop=STOPCHUNK)
        xyz = cmd.get_coords(selection=SELECTION, state=0)
        nstates = cmd.count_states('traj')
        if nstates == 1:
            sys.exit()
        if STOPCHUNK >= STOP:
            delta = STOPCHUNK - STOP
        else:
            delta = None
        if not os.path.exists('%s.pdb'%OUTBASEFILENAME):
            cmd.save('%s.pdb'%OUTBASEFILENAME, selection=SELECTION, state=1)
            print "Number of atoms in selection: %d"%xyz.shape[1]
        nframe = cmd.count_states()
        natom = cmd.count_atoms(SELECTION, state=1)
        xyz = xyz.reshape(nframe, natom, 3)[1:] # The first frame is the pdb file
        if delta is not None:
            xyz = xyz[:-delta]
        f.write(xyz)
        if delta is not None:
            sys.exit()
        START = STOPCHUNK + 1
        cmd.delete('traj')
cmd.quit()

trajselect_mdtraj.py

#!/usr/bin/env python
# -*- coding: UTF8 -*-

# Author: Guillaume Bouvier -- guillaume.bouvier@pasteur.fr
# https://research.pasteur.fr/en/member/guillaume-bouvier/
# 2018-11-19 16:26:10 (UTC+0100)

import mdtraj as md
import sys
import os

try:
    PDBFILENAME = sys.argv[1]
    TRJFILENAME = sys.argv[2]
    SELECTION = sys.argv[3]
    OUTBASEFILENAME = sys.argv[4]

except IndexError:
    print "Help:"
    print "    pdbselect input.pdb input_traj_file 'selection string' output_base_filename"
    print "    The 'selection string' uses MDTraj selection expressions"
    print "    (see: http://mdtraj.org/latest/atom_selection.html)"
    sys.exit()

total_frames = 0
with md.formats.DCDTrajectoryFile('%s.dcd'%OUTBASEFILENAME, mode='w') as f:
    for chunk in md.iterload(TRJFILENAME, top=PDBFILENAME):
        chunk = chunk.atom_slice(chunk.top.select(SELECTION), inplace=True)
        total_frames += chunk.n_frames
        xyz = chunk.xyz
        if not os.path.exists('%s.dcd'%OUTBASEFILENAME):
            chunk[0].save_pdb('%s.pdb'%OUTBASEFILENAME)
            print "• Number of atoms in selection: %d"%xyz.shape[1]
        sys.stdout.write("    ‣ Writting frames: %d\r"%total_frames)
        sys.stdout.flush()
        f.write(xyz)
print "        ‣ %d frames written"%total_frames