Commit 30d084ac authored by David M. Rogers's avatar David M. Rogers
Browse files

Added equivalent (untested) code to store original pdbqt info.

parent 4d2d8b63

helpers.py

+19 −0
Original line number Diff line number Diff line
@@ -2,6 +2,10 @@ import re 
import tarfile as tf

from pathlib import Path



def shard_id(tarname):

    st = tarname.find('output_p')

    return tarname[st+7:].split('.', 1)[0] # e.g. p43812_15



# iterator over tar-file

def tar_iter(name, ext=None):

    inp = Path(name)
@@ -41,6 +45,21 @@ def grep_all(f, *keys): 
                out[k].append(line)

    return out



# parse relevant info from pdbqt

def get_pdbqt_info(f):

    atom = 0

    tors = None

    for line in f:

        try:

            if line[:6] == b"ATOM  " or line[:6] == b"HETATM":

                atom = max(atom, int(line[6:11]))

            elif line[:7] == b"TORSDOF":

                tors = int(line[7:])

        except ValueError:

            print("Error getting info. `{}`".format(line))

            raise

    return atom, tors



def parse_dlg(f):

    #dlg = grep_all(f, "DOCKED:", "Estimated", "Number of atoms", "Number of rotat")

    dlg = grep_all(f, "DOCKED:")

map_fn.py

+2 −1
Original line number Diff line number Diff line
@@ -15,11 +15,12 @@ project = signac.get_project() 


# relies on global variables extension and map_fn (see example run.py)

def insert_loop(tname):

    shard = shard_id(tname)

    with signac.pymongo_buffered(project):

        fname = None

        for fname, name, f in tar_iter(tname, extension):

            job = project.open_job(statepoint = {'real_id':name})

            map_fn(job, f)

            map_fn(job, f, shard)

        if fname is None:

            print("No %s files found in %s."%(extension, tname))

pdbqts.py

0 → 100644
+14 −0
Original line number Diff line number Diff line 
from helpers import *



extension = "pdbqt"



def map_fn(job, f, shard):

    try:

        atoms, tors = get_pdbqt_info(f)

        if tors is None:

            tors = 0

    except:

        return

    job['shard'] = shard

    job['atoms'] = atoms

    job['tors'] = tors

run.py

+2 −2
Original line number Diff line number Diff line
@@ -15,7 +15,7 @@ class Tmp: 
    def __init__(self):

        self.doc = {}



def add_score(job, f):

def add_score(job, f, shard):

    #job = Tmp()



    en, conf = parse_dlg(f)
@@ -44,7 +44,7 @@ def add_score(job, f): 
    job.doc[dist_key]  = np.sqrt(np.sum(dr*dr))



# extension : xvg, pdbqt

# map_fn : (job, file obj) -> ()

# map_fn : (job, file obj, shard_id) -> ()

#     -- can run, e.g. job.doc[attr] = val

# the function to map

extension = "dlg"

run_example.py

+2 −2
Original line number Diff line number Diff line 
#from helpers import *

#import numpy as np



def add_score(job, f):

def add_score(job, f, shard):

    job.doc['score'] = 0.0



# extension : xvg, pdbqt

# map_fn : (job, file obj) -> ()

# map_fn : (job, file obj, shard : str) -> ()

#     -- can run, e.g. job.doc[attr] = val

# the function to map

extension = "dlg"