Wicked Fast Cheminformatics with NVIDIA RAPIDS


Graphics Processing Units (GPUs) have revolutionized scientific computing.  Scientists have been using GPUs to achieve significant speed-ups in fields ranging from molecular dynamics to machine learning.  Unfortunately, programming GPUs is a rather painful process that requires considerable expertise. Fortunately for those of us who'd prefer to forgo the travails of CUDA programming, NVIDIA has released the RAPIDS library, which makes it easy to perform a wide array of data science operations on a GPU.  In this post, I'll present a few examples of how we can use RAPIDS to speed-up a few tasks that we commonly perform in Cheminformatics.  As usual, a Jupyter notebook containing all of the code associated with this post is available on GitHub.

2020 -06-23 I made a couple of changes to the code that slightly changed the runtimes and the trustworthiness values for t-SNE.  The conclusions are the same, RAPIDS ROCKS!

Installation
I've been following RAPIDS since its initial public release and the installation process has gone from incredibly painful to very easy.  Here's how I installed the RDKit and RAPIDS into a conda environment.

1. Install the RDKit into a new conda environment and activate the environment.
conda create -c rdkit -n rdkit_2020_06
conda activate rdkit_2020_06

2.  Check the versions of Python and CUDA.
python --version
nvidia-smi

3. Use the handy RAPIDS configuration tool to set up a conda install command to fit my versions of Python and CUDA.
conda install -c rapidsai -c nvidia -c conda-forge \
    -c defaults rapids=0.14 python=3.6 cudatoolkit=10.1

 If you're running a Jupyter notebook on AWS, you'll probably need to execute this command before starting your notebook for the first time. 
pip install environment_kernels


Introduction
At the heart of RAPIDS is a data structure known as the CUDA dataframe or cudf.   One can think of a cudf as a Pandas dataframe that lives on a GPU.    This functionality is in a library called cudf. It's trivial to create a cudf from a Pandas dataframe.

import cudf
import pandas as pd
df = pd.read_csv("chembl_100k.smi.gz",sep=" ",names=["SMILES","Name"])
cu_df = cudf.from_pandas(df)

In the examples we'll be looking at here, we'll primarily be dealing with chemical fingerprints in a numpy array.  The RAPIDS examples contain a handy function for converting a numpy array to a cudf.

def np2cudf(arr):
    # convert numpy array to cudf dataframe
    df = pd.DataFrame({'fea%d'%i:arr[:,i] for i in range(arr.shape[1])})
    pdf = cudf.DataFrame()
    for c,column in enumerate(df):
        pdf[str(c)] = df[column]
    return pdf

 Note that this isn't the most memory-efficient way to hold fingerprints in memory.  If you're dealing with millions of molecules, you may run out of GPU memory. 

In addition to speeding up operations such as searching and sorting dataframes, RAPIDS contains optimized implementations of several machine learning algorithms that operate on a cudf.  This functionality is in a Python library called cuml.   In this post, we'll take a look at GPU implementations of a few unsupervised machine learning algorithms that I've covered in previous posts.
K-Means Clustering
As I discussed in my previous post, K-Means clustering provides an efficient means of dividing a set of molecules into a predefined number of clusters.  We'll benchmark the performance of the RAPIDS implementation of k-means by comparing it with the MiniBatch K-Means implementation in scikit-learn.  All of the benchmarks discussed here were run on a p3.2xlarge instance on AWS with an NVIDA V100 GPU.  In the plot below, we compare the time required to cluster 10,000 molecules into 100 to 1,000 clusters.  
As we can see from the plot above, the RAPIDS GPU implementation is considerably faster than the CPU implementation in scikit-learn.  In order to assess the performance differences, we'll plot the ratio of CPU runtime to GPU runtime.   From the plot below we can see that the RAPIDS GPU implementation is 5 to 45 times faster than the sklearn CPU implementation.  
Finally, to get a better idea of performance, we'll run a benchmark with a larger set of 100,000 molecules.  We see that we can cluster 100K fingerprints into 1,000 clusters in less than 8 seconds.  I'm not sure why the run with 700 clusters was slower but I ran this a few times and saw the same behavior.  Note that, for comparison, I tried to run this clustering on a CPU but the clustering was taking hours and I was getting bored.  Datasets of this size should definitely be run on a GPU. 

Visualizing Chemical Space with t-SNE and UMAP
In a previous post, I wrote about t-SNE a dimensionality reduction technique that can be used to visualize the chemical space covered by a set of molecules.  A couple of comments on that post suggested UMAP as an alternative to t-SNE.  Fortunately, both of these algorithms are implemented in RAPIDS.  

As a test, I used t-SNE to map the chemical space for a set of 1,495 drugs from the ChEMBL database.  As above, the RAPIDS implementation was quite a bit faster.  Mapping the 1,495 molecules took 6.7 seconds using the CPU implementation in scikit-learn while the GPU implementation in RAPIDS only took 1 second.  While it's nice that RAPIDS was faster the results don't look the same as those obtained with scikit-learn.  


In order to further test this, I calculated the trustworthiness, a measure of how faithfully the distances in the low dimensional space reproduce the corresponding distances in the high dimensional space.  This value ranges between 0 and 1, with higher values indicating that the projection more faithfully reproduces the distances in the original high dimensional space.  As you can see below, the scikit-learn implementation appears to have a slightly higher value of trustworthiness than the RAPIDS implementation. 


While the RAPIDS implementation of t-SNE is quick,  I need to understand it better before integrating it into my standard workflow. 

I performed a similar comparison on the same ChEMBL drug dataset with a CPU implementation of UMAP and the GPU implementation in RAPIDS.  In this case, the resulting projections and the trustworthiness values were similar.   

Note that we shouldn't compare the trustworthiness between t-SNE and UMAP.  The t-SNE results were calculated on the first 50 principal components of the fingerprints, while UMAP was calculated based on the full 2048 bit fingerprints. 


The fact that the RAPIDS implementation of UMAP is almost 20x faster opens a lot of possibilities.  Among other things, this sort of speed makes it practical to examine the impact of parameter settings on the output of a method. For instance, in UMAP one can specify the number of neighbor molecules used to calculate the map.  Since the RAPIDS implementation is fast, these scans can be performed in seconds.  Here I evaluate the relationship between trustworthiness and the number of neighbors and show that it maxes out with three neighbors. 

We can also scan another parameter, the minimum distance between clusters, and show that it has a minimal impact on UMAP performance. 


Conclusion
In this post, we've looked at a couple of unsupervised machine learning methods that have fast GPU implementations in NVIDIA's RAPIDS package.  The functions in RAPIDS are very easy to use and can provide a simple means of speeding up your Cheminformatics workflows.  However, as we saw with t-SNE, the results can sometimes be less than satisfactory.  It's always good to compare with an established, gold-standard implementation to ensure that your getting the results you expected.  There's a lot more to RAPIDS and I'd urge you to try out some of the demos and adapt them to your work. I'll be covering more on RAPIDS in future posts. 





Comments

Post a Comment

Popular posts from this blog

AI in Drug Discovery 2023 - A Highly Opinionated Literature Review (Part I)

Image
Here’s the first part of my review of some interesting machine learning (ML) papers I read in 2023.  As with the previous editions , this shouldn’t be considered a comprehensive review.  The papers covered here reflect my research interests and biases, and I’ve certainly overlooked areas that others consider vital.  This post is pretty long, so I've split it into three parts, with parts II and III to be posted in the next couple of weeks.    I. Docking, protein structure prediction, and benchmarking II. Large Language Models, active learning, federated learning, generative models, and explainable AI III. Review articles 2023 was a bit of a mixed bag for AI in drug discovery.  Several groups reported that the deep learning methods for protein-ligand docking weren’t quite what they were initially cracked up to be.  AlphaFold2 became pervasive, and people started to investigate, with mixed success, the utility of predicted protein structures.  There were reports of significant advanc
Read more

Generative Molecular Design Isn't As Easy As People Make It Look

Image
I was taken aback by a recent CNBC article entitled “ Generative AI will be designing new drugs all on its own in the near future ”.  I should know better than to pay attention to AI articles in the popular press, but I feel that even scientists working in drug discovery may have a skewed perception of what generative AI can and can’t do.  To understand exactly what’s involved, it might be instructive to walk through a typical generative molecular design workflow and point out a few things.  First, these programs are far from autonomous.  Even when presented with a well-defined problem, generative algorithms produce a tremendous amount of nonsense.  Second, domain expertise is essential when sifting through the molecules produced by a generative algorithm.  Without a significant medicinal chemistry background, one can’t make sense of the results.  Third, while a few nuggets exist in the generative modeling output, a lot of work and good old-fashioned cheminformatics are required to ext
Read more

AI in Drug Discovery - A Highly Opinionated Literature Review (Part II)

Image
Picking up where we left off in Part I , this post covers several other ML in drug discovery topics that interested me in 2023.  Some areas, like large language models, are new, and most of the work is at the proof-of-concept stage.  Others, like active learning, are more mature, and several groups are starting to explore nuances of the methods.   Here’s the structure of Part II.  4. Large Language Models 5. Active Learning 6. Federated Learning 7. Generative Models 8. Explainable AI 9. Other Stuff 4. Large Language Models The emergence of GPT-4 and ChatGPT brought considerable attention to large language models (LLMs) in 2023.  In November and December, several large pharmas held “AI Day” presentations featuring LLM applications for clinical trial data analysis. Many of these groups demonstrated the ability of LLMs to ingest large bodies of unstructured clinical data and subsequently generate tables and reports based on natural language queries.  Aside from some very brief demos on co
Read more