volumeanalysis_parallel.py

#!/usr/bin/env python
import argparse
import gencavitylib
import numpy
import os
import subprocess
import tempfile
import multiprocessing
import types
import copy_reg

def _pickle_method(m):
    if m.im_self is None:
        return getattr, (m.im_class, m.im_func.func_name)
    else:
        return getattr, (m.im_self, m.im_func.func_name)

copy_reg.pickle(types.MethodType, _pickle_method)
        

class VolumeAnalysis(object):
    expecttemplate = '''\
#!/usr/local/bin/expect
spawn /gscratch3/lchong/ajd98/apps/voidoo_linux/lx_voidoo
expect "Type of calculation (C/V/R/Q)         ? (C)"
send "V\r"
expect "Do you want extensive output          ? (N)"
send "\r"
expect "Library file ? (cavity.lib)"
send "%LIBRARYFILE%\r"
expect "PDB file name ? (in.pdb)"
send "%PDBFILE%\r"
expect "Primary grid spacing (A) ? (   1.000)"
send "\r"
expect "Probe radius (1.4 A for water) ? (   0.000) "
send "\r"
expect " Nr of volume-refinement cycles        ? (         10) "
send "30\r"
expect "Grid-shrink factor                    ? (   0.900) "
send "\r"
expect " Convergence criterion (A3)            ? (   0.100)"
send ".5\r"
expect " Convergence criterion (%)             ? (   0.100) "
send ".001\r"
expect "Create protein-surface plot file      ? (N) "
send "\r"
expect "***** VOIDOO ***** VOIDOO ***** VOIDOO ***** VOIDOO ***** VOIDOO ***** "
send "\r"
'''
    nprocesses = 16

    def __init__(self):
        self._parse_args()
        self.generate_cavity_lib()
        self.volumes = []
        self.go()

    def _parse_args(self):
        parser = argparse.ArgumentParser()
        parser.add_argument('--maindir', dest='maindir',
                            help='The main simulation directory '
                                 '(e.g., ff14sb.xray.1).',
                            required=True)
        parser.add_argument('--topology', dest='topologypath',
                            help='The path to the topology files',
                            required=True)
        parser.add_argument('--outdir', dest='outdirpath',
                            help='The path to the output directory.',
                            required=True)

        self.args = parser.parse_args()

    def get_simname(self):
        simname = self.args.maindir.split('/')[-1]
        return simname

    def generate_cavity_lib(self):
        self.cavitylibpath = os.path.join(self.args.outdirpath, 
                                          'cavity.lib.autogen')

        gencavitylib.CavityLibGen(self.args.topologypath, self.cavitylibpath)
        
        # Now read the file we just wrote, and remove the "RESI" entries for
        # residues that should not be included in the volume calculation
        with open(self.cavitylibpath,'r') as f:
            lines = f.readlines()
        newlines = []
        for line in lines:
            if ('WAT' in line) or ('Cl-' in line) or ('Cl' in line):
                pass
            elif 'nan' in line:
                newlines.append(line.replace('nan','0.000'))
            else:
                newlines.append(line) 

        # Write the filtered list back to the file.
        with open(self.cavitylibpath, 'w') as f:
            f.writelines(newlines)

    def runvoidoo(self, pdbfilename, expectfilename, tpid):
        '''
        Run voidoo on the pdb file with path
            "pdbfile.name" + ".pdb." + str(tpid)
        using the named temporary file name "expectfilename" for the expect script that
        controls voidoo
        '''
        expectcmd = self.expecttemplate
        expectcmd = expectcmd.replace('%LIBRARYFILE%',self.cavitylibpath)
        pdbpath = "{:s}.pdb.{:d}".format(pdbfilename, tpid)
        expectcmd = expectcmd.replace('%PDBFILE%', pdbpath)
        with open(expectfilename, 'w+') as expectfile:
            expectfile.write(expectcmd)
            expectfile.flush()
            process = subprocess.Popen(['expect', expectfile.name], 
                                       stdout=subprocess.PIPE)
            process.wait()
            stdout, stderr = process.communicate()
            
            # Finally, parse the output from VOIDOO
            s = [line for line in stdout.split('\n') if ("Protein volume (A3)" in line)][-1]
            volume = float(s.split()[5][:-1])
            
            # Reset the expect script (delete its contents)
            expectfile.seek(0)
            #expectfile.file.truncate()
            expectfile.truncate()
        # Delete the pdb file
        os.remove(pdbpath)

        return volume

    def runvoidoowrapper(self, args):
        '''
        args should be a tuple that may be passed to self.runvoidoo as args.
        '''
        return self.runvoidoo(*args)
        
    def go(self):
        print("Using {:d} processes".format(self.nprocesses))
        pool = multiprocessing.Pool(self.nprocesses)

        expectfiles = [tempfile.NamedTemporaryFile() \
                       for i in range(100)]
        for segid in range(1,10001):
            print("Working on segment {:d}".format(segid))
            # write a cpptraj script to write out the frames as separate pdb files
            cpptrajscript = tempfile.NamedTemporaryFile()
            pdbfile = tempfile.NamedTemporaryFile()
            pdbfilename = pdbfile.name
            pdbfile.close()


            cpptrajscript.write("parm {:s}\n".format(self.args.topologypath))
            if "ff14sb.xray.1" in self.args.maindir or "ff14sb.nmr.1" in self.args.maindir:
                cpptrajscript.write("trajin {:s}/{:04d}/{:04d}_cut.nc\n"\
                                    .format(self.args.maindir, segid, segid))  
            else:
                cpptrajscript.write("trajin {:s}/{:05d}/{:05d}_cut.nc\n"\
                                    .format(self.args.maindir, segid, segid))  
            cpptrajscript.write("strip :WAT,Cl,Cl-,SOL\n")
            cpptrajscript.write("trajout {:s}.pdb pdb multi nobox\n".format(pdbfilename))
            cpptrajscript.write("go\n")
            #cpptrajscript.file.flush()
            cpptrajscript.flush()

            # Call cpptraj to write out all the pdb files.
            subprocess.Popen(['/ihome/lchong/ajd98/apps/amber14/bin/cpptraj', '-i', 
                              cpptrajscript.name], stdout=subprocess.PIPE).wait()
             
            
            # Now iterate over the 100 pdb files
            mappingargs = [(pdbfilename, expectfiles[i].name, i+1) for i in range(100)]
            volumes = pool.map(self.runvoidoowrapper, mappingargs)
            self.volumes += volumes

        outpath = os.path.join(self.args.outdirpath, 'volume')
        numpy.save(outpath, numpy.array(self.volumes))

        
      
if __name__ == '__main__':
    VolumeAnalysis()