RAD (Retrieval Augmented Docking)

RAD is a scalable virtual screening library using HNSW graphs and distributed computing. The architecture supports deployment from single machines to HPC clusters using a central coordination service.

Requirements

• Redis
• Python >=3.11
• GCC >= 9.3

Installation

git clone --recursive https://github.com/keiserlab/rad.git
cd rad
pip install .

We also provide a Dockerfile containing all required software.

Architecture Overview

RAD uses a service-oriented design with three main components:

1. HNSW Service: Handles HNSW neighbor searches and SMILES lookup
2. Coordination Service: Manages work distribution, acts as HNSW proxy, and maintains state via Redis
3. Distributed Workers: Lightweight scoring processes that can run anywhere with only Redis access

Running RAD

Basic Workflow

1. Build HNSW graph from molecular fingerprints
2. Create SQLite database mapping node keys to SMILES
3. Define a SMILES-based scoring function
4. Initialize RAD services and run traversal

Constructing the HNSW

Constructing the HNSW graph consists of setting the construction parameters expansion_add and connectivity and then adding each molecule by providing a numerical key and its fingerprint.

expansion_add controls the number of candidates considered as potential neighbors during element insertion, while connectivity controls how many of these candidates are actually connected to the inserted element.

from usearch.index import Index
hnsw = Index(
ndim = 1024, # 1024 bit fingerprint
dtype='b1', # For packed binary fingerprints
metric='tanimoto',
connectivity = 8,
expansion_add = 400
)
fingerprints = ...
keys = np.arange(len(fingerprints))
hnsw.add(keys, fingerprints, log="Building HNSW")

The fingerprints are expected to be an (n x d/8) numpy array where each row is a packed binary fingerprint. e.g turning a 1024-bit binary fingerprint into a 128 uint8 fingerprint with np.packbits(). See the example notebook for more details.

Creating SQLite Database for SMILES mapping

RAD integrates with SQLite to provide SMILES directly to scoring functions:

importsqlite3# Create database mapping HNSW keys to SMILESconn=sqlite3.connect('molecules.db')
cursor=conn.cursor()
cursor.execute(""" CREATE TABLE nodes ( node_key INTEGER PRIMARY KEY, smi TEXT NOT NULL )""")
# Insert SMILES dataforkey, smilesinzip(keys, smiles):
cursor.execute("INSERT INTO nodes (node_key, smi) VALUES (?, ?)", (key, smiles))
cursor.execute("CREATE INDEX idx_nodes_node_key ON nodes(node_key)")
conn.commit()
conn.close()

Defining a SMILES-Based Scoring Function

Scoring functions receive SMILES strings and return a score. Numerically smaller scores are considered better. Here is a mock example:

defscore_fn(smiles: str) ->float:
score=calculate_docking_score(smiles)
returnscore# Lower scores are better

Initializing RAD Services

With the HNSW index, SMILES database, and scoring function ready, initialize the RAD traverser:

fromrad.hnsw_serviceimportcreate_local_hnsw_servicefromrad.traverserimportRADTraverser# Create HNSW service with database integrationhnsw_service=create_local_hnsw_service(hnsw, database_path='molecules.db')
# Create traverser with SMILES-based scoringtraverser=RADTraverser(hnsw_service=hnsw_service, scoring_fn=score_fn)

Deployment Modes

Local Deployment (single machine):

traverser=RADTraverser(hnsw_service=hnsw_service, scoring_fn=score_fn)

Distributed Deployment (HPC):

traverser=RADTraverser(
hnsw_service=hnsw_service, scoring_fn=score_fn,
redis_host='head-node.cluster',
redis_port=6379,
namespace='job_12345'
)

Remote HNSW Service:

fromrad.hnsw_serviceimportcreate_remote_hnsw_service# Start HNSW server elsewhere OR use the publicly provided server# python scripts/start_hnsw_server.py --database-path molecules.db --hnsw-path index.usearch --port 8000hnsw_service=create_remote_hnsw_service("https://rad.docking.org")
traverser=RADTraverser(hnsw_service=hnsw_service, scoring_fn=score_fn)

Priming the RAD Traverser

The traverser is 'primed' by finding and scoring the nodes on the top layer of the HNSW graph and initializing the priority queue. This should only be run once.

traverser.prime()

Performing the traversal

The traversal proceeds until a maximum number of molecules is scored or a timeout is reached:

# Run traversal until 100k molecules are scoredtraverser.traverse(n_workers=4, n_to_score=100_000)
# Or run traversal for a specific timetraverser.traverse(n_workers=4, timeout=3600) # 1 hour

Accessing the results

RAD provides two methods for accessing results:

Traversal Order:

# Get molecules in the order they were discoveredmolecules=traverser.get_molecules() # All moleculesfirst_100=traverser.get_molecules(100) # First 100 moleculesfornode_id, score, smilesinmolecules:
print(f"Node {node_id}: {smiles} (score: {score})")

Best Molecules:

# Get top-scoring molecules regardless of discovery orderbest_molecules=traverser.get_best_molecules(10) # Top 10 by scorefornode_id, score, smilesinbest_molecules:
print(f"Top hit: {smiles} (score: {score})")

Service Management and Cleanup

Gracefully shutdown all services:

traverser.shutdown()

Advanced Usage

Starting HNSW Server Independently:

# Start dedicated HNSW server with database
python scripts/start_hnsw_server.py \
--hnsw-path /data/index.usearch \
--database-path /data/molecules.db \
--host 0.0.0.0 \
--port 8000

Example Usage

The examples/ folder contains a Jupyter notebook demonstrating the construction and traversal of the DUDE-Z DOCK HNSW investigated in the original RAD paper.

For a larger billion-scale application and integration with Chemprop see the repo at https://github.com/bwhall61/lsd

References

The original HNSW paper by Yury Malkov and Dmitry Yashunin.

The original RAD paper by Brendan Hall and Michael Keiser.

The lsd.docking.org paper by Brendan Hall, Tia Tummino, et al. shows a billion-scale application and integration with Chemprop ML models.

And then most importantly, the HNSW graph code is built on the usearch library so large thanks to Ash Vardanian for his awesome HNSW library!