Clustering 2.1 Million Compounds for $5 With a Little Help From Amazon & Facebook

In this post, I'll talk about how we can use FAISS, an Open Source Library from Facebook, to cluster a large set of chemical structures.   As usual, the code associated with this post is on GitHub. 

As I wrote in a previous post, K-means clustering can be a useful tool when you want to partition a dataset into a predetermined number of clusters.   While there are a number of tricks for speeding up k-means (also mentioned in the previous post), it can still take a long time to run when the number of clusters, or the number of items being clustered, is large.

 One of the great things about Cheminformatics these days is that we can take advantage of advances in other fields.  One such advance comes in the form of a software library called FAISS that was released as Open Source by Facebook.  FAISS is a library of routines for performing extremely efficient similarity searches.  While many of us think about similarity searches in terms of chemical fingerprints, similar techniques are at the heart of modern commerce.   Our characteristics, such as likes, dislikes, and purchases are represented as vectors.  Similarity searches among these vectors enable websites to tailor content and point out that people who bought this also bought that.  Of course, Facebook would like to do this quickly, thus the motivation behind FAISS.

 In addition to similarity search, FAISS also provides routines to perform k-means clustering.  As implemented, FAISS can run on either a CPU or a GPU.  Having seen the performance increases that can be achieved on a GPU, I was initially drawn to the GPU implementation.  Unfortunately, I found it difficult to get the data I was using into GPU memory.  This is something I'll spend some more time with in the future, but for now, I have work to do.

 Being lazy pragmatic, I decided to try out the CPU implementation.  Fortunately, FAISS scales across multiple CPUs and CPU memory is a lot easier to come by.  On AWS, I can rent an m5.24xlarge instance with 96 CPUs for $4.61/hr.  That's a lot of CPU!  Let's see what we can do with it. 

Ok, Here's the problem we want to solve.  We have a commercial library of 2.1 million molecules, and we would like to cluster this to 50K clusters.  We can break this problem down into two steps.  The first step is to generate the fingerprints and write them to disk.

Step 1 - Creating a Fingerprint File 

 In a previous post, I used the parquet library to store a Pandas dataframe containing fingerprints to disk.  After some experimentation, I found that this approach was pretty slow when the dataset got bigger.  For datasets containing more than a million fingerprints, I found that storing the data in an HDF5 file was much faster, the pyhd5 library makes this pretty easy.  The other nice aspect to the HDF5 files is that it's easy to integrate other information like SMILES strings and molecule names into the same file. 

The code in the GitHub repo has a python script called gen_fp.py which will generate fingerprints and write them to an HDF5 file.   I just went with a simple 256 bit Morgan fingerprint as the default.  It's easy enough to change this if you want to use something else.  Using the script is pretty simple.

gen_fp.py test2100K.smi test2100K.h5

It took me 11.9 minutes to generate fingerprints for 2.1 million molecules on a single core.

Step 2 - Clustering 2.1 million molecules into 50,000 clusters

The GitHub repo has another script called faiss_kmeans.py that performs the clustering.   Again, the command line is pretty simple.

faiss_kmeans.py test2100K.h5 50000

This clustering will produce two output files.

  • centers.smi contains the smiles for the most central molecule in each cluster
  • detail.csv contains the SMILES, molecule name, distance to the cluster center, and cluster number for each molecule in the input file.
Here's a screenshot of htop running with 96 cores clustering with FAISS, it's a beautiful thing! 



Clustering 2.1 million molecules took 22 minutes on 92 cores, not too bad.  It's amazing what you can do with on-demand computing and a little bit of open source code. 

I haven't spent a lot of time pushing the limits of this method, but a lot more may be possible, stay tuned. 

Coda 

A few questions came up on Twitter so I thought I'd answer them here. 

1. Why are you doing this? 

As I mentioned in my previous blog post on K-means, this is part of my workflow when I select compounds for a screening collection.   I want a diverse set of compounds with lots of analogs that I can purchase.  My workflow typically goes like this.

  • Filter to remove molecules with objectionable properties or functionality.  See my post on filtering chemical libraries for more information. 
  • Cluster the remaining molecules using K-means
  • Select the molecule from each cluster with the largest number of commercially available analogs. 
2.  What is the size of the largest/smallest clusters? Or what does the distribution of cluster sizes look like?

The largest cluster has 296 members, the smallest has 1.  The arithmetic mean is 42 and the median is 36.   The distribution looks like this. 


3. Did you like what you saw, ie did the clusters look sensible?

I can't say that I looked at all or even most of the clusters, but the ones I looked at made sense. 


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