Prepare antibody data for Pafnucy

prepare.py

import numpy as np
import pandas as pd
import h5py

import pybel
from tfbio.data import Featurizer

import os


def input_file(path):
    """Check if input file exists."""

    path = os.path.abspath(path)
    if not os.path.exists(path):
        raise IOError('File %s does not exist.' % path)
    return path


def output_file(path):
    """Check if output file can be created."""

    path = os.path.abspath(path)
    dirname = os.path.dirname(path)

    if not os.access(dirname, os.W_OK):
        raise IOError('File %s cannot be created (check your permissions).'
                      % path)
    return path


def string_bool(s):
    s = s.lower()
    if s in ['true', 't', '1', 'yes', 'y']:
        return True
    elif s in ['false', 'f', '0', 'no', 'n']:
        return False
    else:
        raise IOError('%s cannot be interpreted as a boolean' % s)


import argparse
parser = argparse.ArgumentParser(
    description='Prepare molecular data for the network',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    epilog='''This script reads the structures of ligands and pocket(s),
    prepares them for the neural network and saves in a HDF file.
    It also saves affinity values as attributes, if they are provided.
    You can either specify a separate pocket for each ligand or a single
    pocket that will be used for all ligands. We assume that your structures
    are fully prepared.\n\n

    Note that this scripts produces standard data representation for our network
    and saves all required data to predict affinity for each molecular complex.
    If some part of your data can be shared between multiple complexes
    (e.g. you use a single structure for the pocket), you can store the data
    more efficiently. To prepare the data manually use functions defined in
    tfbio.data module.
    '''
)

parser.add_argument('--ligand', '-l', required=True, type=input_file, nargs='+',
                    help='files with ligands\' structures')
parser.add_argument('--ligand_format', type=str, default='mol2',
                    help='file format for the ligand,'
                         ' must be supported by openbabel')
parser.add_argument('--output', '-o', default='./complexes.hdf',
                    type=output_file,
                    help='name for the file with the prepared structures')
parser.add_argument('--mode', '-m', default='w',
                    type=str, choices=['r+', 'w', 'w-', 'x', 'a'],
                    help='mode for the output file (see h5py documentation)')
parser.add_argument('--affinities', '-a', default=None, type=input_file,
                    help='CSV table with affinity values.'
                         ' It must contain two columns: `name` which must be'
                         ' equal to ligand\'s file name without extenstion,'
                         ' and `affinity` which must contain floats')
parser.add_argument('--verbose', '-v', default=True, type=string_bool,
                    help='whether to print messages')

args = parser.parse_args()

num_ligands = len(args.ligand)
if args.verbose:
    print('%s molecules to prepare:' % num_ligands)
    for ligand_file in args.ligand:
            print(' molecule: %s' % ligand_file)
    print('\n\n')


if args.affinities is not None:
    affinities = pd.read_csv(args.affinities)
    if '-logKd' not in affinities.columns:
        raise ValueError('There is no `-logKd` column in the table')
    elif 'pdbid' not in affinities.columns:
        raise ValueError('There is no `pdbid` column in the table')
    affinities = affinities.set_index('pdbid')['-logKd']
else:
    affinities = None

featurizer = Featurizer()


with h5py.File(args.output, args.mode) as f:
    for ligand_file in args.ligand:
        # use filename without extension as dataset name
        name = os.path.splitext(os.path.split(ligand_file)[1])[0]

        if args.verbose:
            print('reading %s' % ligand_file)
        try:
            ligand = next(pybel.readfile(args.ligand_format, ligand_file))
        except:
            raise IOError('Cannot read %s file' % ligand_file)

        ligand_coords, ligand_features = featurizer.get_features(ligand, molcode=-1)

        centroid = ligand_coords.mean(axis=0)
        ligand_coords -= centroid

        data = np.concatenate(
            (np.concatenate((ligand_coords,)),
             np.concatenate((ligand_features,))),
            axis=1,
        )

        dataset = f.create_dataset(name, data=data, shape=data.shape,
                                   dtype='float32', compression='lzf')
        if affinities is not None:
            dataset.attrs['affinity'] = affinities.loc[name]
if args.verbose:
    print('\n\ncreated %s with %s structures' % (args.output, num_ligands))