get_hdf5_program.py

import csv
from csv import reader
import argparse
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem import Draw
from scipy.sparse import csr_matrix
import h5sparse
import numpy

#PARTS 2 + 3

def get_hdf5_from_tfp(csv, ofn, dimension = 1024, dtype = numpy.bool):

    indptr = [0]
    indices = []
    data = []

    num_mols = 0

    with open(csv, 'r') as mols:
        csv_reader = reader(mols)
        for row in csv_reader:
            smile = row[1]
            m = Chem.MolFromSmiles(smile)
            m.SetProp("_Name", row[0])

            fp = Chem.rdmolops.RDKFingerprint(m, fpSize = dimension)

            for i in range(dimension):
                if fp.GetBit(i) is True:
                    indices.append(i)
                    data.append(1)

            indptr.append(len(indices))
            num_mols += 1

    mol_fps = csr_matrix((data, indices, indptr), shape = (num_mols, dimension), dtype = dtype)
    #print(mol_fps)
    print('the dimension of the returned sparse matrix: %d*%d' % mol_fps.shape)

    f = h5sparse.File(ofn + '.hdf5', 'w')
    f.create_dataset('mol_fps', data = mol_fps)

    f.close()

    return mol_fps

def main():
    parser = argparse.ArgumentParser(formatter_class = argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument(
        '--csv',
        type = str)
    parser.add_argument(
        '--ofn',
        type = str)
    parser.add_argument(
        '--dim',
        type = int,
        default = 1024)
    parser.add_argument(
        '--dtype',
        default = numpy.bool)

    args = parser.parse_args()

    get_hdf5_from_tfp(args.csv, args.ofn, args.dim, args.dtype)