cheminformatics_benchmark.py

__author__ = 'zero323'

from rdkit import DataStructs
from rdkit.Chem.Fingerprints import FingerprintMols
from rdkit import Chem
from indigo import Indigo
import pybel
import time
import cPickle
import numpy as np

INPUT_SMILES_FILE = 'smiles.smi'
MESSAGE = '{0}\t{1}\t{2}'
QUERY_SIZE = 1000
NTIMES = 1
molecule_strings = map(lambda x: x.strip(), open(INPUT_SMILES_FILE).readlines())


def test_indigo():
    indigo = Indigo()
    package_name = 'indigo'
    np.random.shuffle(molecule_strings)

    time_start = time.clock()
    molecules = map(lambda molecule_string: indigo.loadMolecule(molecule_string), molecule_strings)
    time_stop = time.clock()
    print(MESSAGE.format(package_name, 'molecule_loading', time_stop - time_start))

    pck_molecules = map(lambda molecule: molecule.serialize(), molecules)

    time_start = time.clock()
    molecules = map(lambda molecule: indigo.unserialize(molecule), pck_molecules)
    time_stop = time.clock()
    print(MESSAGE.format(package_name, 'molecule_deserialize', time_stop - time_start))

    time_start = time.clock()
    fingerprints = map(lambda molecule: molecule.fingerprint("sim"), molecules)
    time_stop = time.clock()
    print(MESSAGE.format(package_name, 'calculating_fingerprints', time_stop - time_start))

    query = fingerprints[:QUERY_SIZE]

    time_start = time.clock()
    for q in query:
        map(lambda f:  indigo.similarity(f, q, 'tanimoto'),  fingerprints)
    time_stop = time.clock()
    print(MESSAGE.format(package_name, 'calculating_similarity', time_stop - time_start))


def test_pybel():
    package_name = 'pybel'
    np.random.shuffle(molecule_strings)

    time_start = time.clock()
    molecules = map(lambda molecule_string: pybel.readstring("smi", molecule_string), molecule_strings)
    time_stop = time.clock()
    print(MESSAGE.format(package_name, 'molecule_loading', time_stop - time_start))

    time_start = time.clock()
    fingerprints = map(lambda molecule: molecule.calcfp(), molecules)
    time_stop = time.clock()
    print(MESSAGE.format(package_name, 'calculating_fingerprints', time_stop - time_start))

    query = fingerprints[:QUERY_SIZE]

    time_start = time.clock()
    for q in query:
        map(lambda f:  q | f,  fingerprints)
    time_stop = time.clock()
    print(MESSAGE.format(package_name, 'calculating_similarity', time_stop - time_start))


def test_rdkit():
    package_name = 'rdkit'
    np.random.shuffle(molecule_strings)
    time_start = time.clock()
    molecules = map(lambda molecule_string: Chem.MolFromSmiles(molecule_string), molecule_strings)
    time_stop = time.clock()
    print(MESSAGE.format(package_name, 'molecule_loading', time_stop - time_start))

    pck_molecules = map(lambda molecule: cPickle.dumps(molecule), molecules)

    time_start = time.clock()
    molecules = map(lambda molecule: cPickle.loads(molecule), pck_molecules)
    time_stop = time.clock()
    print(MESSAGE.format(package_name, 'molecule_deserialize', time_stop - time_start))

    time_start = time.clock()
    fingerprints = map(lambda molecule: FingerprintMols.FingerprintMol(molecule), molecules)
    time_stop = time.clock()
    print(MESSAGE.format(package_name, 'calculating_fingerprints', time_stop - time_start))

    pck_fingerprints = map(lambda fingerprint: cPickle.dumps(fingerprint), fingerprints)

    time_start = time.clock()
    fingerprints = map(lambda fingerprint: cPickle.loads(fingerprint), pck_fingerprints)
    time_stop = time.clock()
    print(MESSAGE.format(package_name, 'fingerprints_deserialize', time_stop - time_start))

    query = fingerprints[:QUERY_SIZE]

    time_start = time.clock()
    for q in query:
        map(lambda f: DataStructs.FingerprintSimilarity(f, q),  fingerprints)
    time_stop = time.clock()
    print(MESSAGE.format(package_name, 'calculating_similarity', time_stop - time_start))


def main(ntimes=NTIMES):
    print('# Test set size: {0} '.format(len(molecule_strings)))
    print('# Query set size: {0} '.format(QUERY_SIZE))
    for i in range(ntimes):
        test_rdkit()
        test_indigo()
        test_pybel()


if __name__ == '__main__':
    main()