Molecular docking via DC-QAOA

Drugs often work by binding to an active site of a protein, inhibiting or activating its function for some therapeutic purpose. Finding new candidate drugs is extremely difficult. The study of molecular docking helps guide this search and involves the prediction of how strongly a certain ligand (drug) will bind to its target (usually a protein).

One of the primary challenges to molecular docking arises from the many geometric degrees of freedom present in proteins and ligands, making it difficult to predict the optimal orientation and assess if the drug is a good candidate or not. One solution is to formulate the problem as a mathematical optimization problem where the optimal solution corresponds to the most likely ligand-protein configuration. This optimization problem can be solved on a quantum computer using methods like the Quantum Approximate Optimization Algorithm (QAOA). This tutorial demonstrates how this paper used digitized-counteradiabatic (DC) QAOA to study molecular docking. This tutorial assumes you have an understanding of QAOA, if not, please see the CUDA-Q MaxCut tutorial found here

The next section provides more detail on the problem setup followed by CUDA-Q implementations below.

Setting up the Molecular Docking Problem

The figure from the paper provides a helpful diagram for understanding the workflow.

docking

There are 6 key steps: 1. The experimental protein and ligand structures are determined and used to select pharmacores, or an important chemical group that will govern the chemical interactions, 2. T wo labeled distance graphs (LAGs) of size \(N\) and \(M\) represent the protein and the ligand, respectively. Each node corresponds to a pharmacore and each edge weight corresponds to the distance between pharmacores. 3. A \(M*N\) node binding interaction graph (BIG) is created from the LAGs. Each node in the BIG graph corresponds to a pair of pharmacores, one from the ligand and the other from the protein. The existence of edges between nodes in the BIG graph are determined from the LAGs and correspond to interactions that can feesibly coexist. Therefore, cliques in the graph correspond to mutually possible interactions. 4. The problem is mapped to a QAOA circuit and corresponding Hamiltonian, and the ground state solution is determined. 5. The ground state will produce the maximum weighted clique which corresponds to the best (most strongly bound) orientation of the ligand and protein. 6. The predicted docking structure is interpreted from the QAOA result and is used for further analysis.

CUDA-Q Implementation

First, the appropriate libraries are imported and the nvidia backend is selected to run on GPUs if available.

[1]:
import cudaq
from cudaq import spin
import numpy as np

# cudaq.set_target('nvidia')

The block below defines two of the BIG data sets from the paper. The first is a smaller example, but it can be swapped with the commented out example below at your discretion. The weights are specified for each node based on the nature of the ligand and protein pharmacores represented by the node

[2]:
# The two graphs input from the paper

# BIG 1

nodes = [0, 1, 2, 3, 4, 5]
qubit_num = len(nodes)
edges = [[0, 1], [0, 2], [0, 4], [0, 5], [1, 2], [1, 3], [1, 5], [2, 3], [2, 4],
         [3, 4], [3, 5], [4, 5]]
non_edges = [
    [u, v] for u in nodes for v in nodes if u < v and [u, v] not in edges
]

print('Edges: ', edges)
print('Non-Edges: ', non_edges)

weights = [0.6686, 0.6686, 0.6686, 0.1453, 0.1453, 0.1453]
penalty = 6.0
num_layers = 3

# BIG 2 (More expensive simulation)
#nodes=[0,1,2,3,4,5,6,7]
#qubit_num=len(nodes)
#edges=[[0,1],[0,2],[0,5],[0,6],[0,7],[1,2],[1,4],[1,6],[1,7],[2,4],[2,5],[2,7],[3,4],[3,5],[3,6],\
#    [4,5],[4,6],[5,6]]
#non_edges=[[u,v] for u in nodes for v in nodes if u<v and [u,v] not in edges]
#print('Edges: ', edges)
#print('Non-edges: ', non_edges)
#weights=[0.6686,0.6686,0.6886,0.1091,0.0770,0.0770,0.0770,0.0770]
#penalty=8.0
#num_layers=8
Edges:  [[0, 1], [0, 2], [0, 4], [0, 5], [1, 2], [1, 3], [1, 5], [2, 3], [2, 4], [3, 4], [3, 5], [4, 5]]
Non-Edges:  [[0, 3], [1, 4], [2, 5]]

Next, the Hamiltonian is constructed.

\[H = \frac{1}{2}\sum_{i \in V}w_i(\sigma^z_i - 1) + \frac{P}{4} \sum_{(i,j) \notin E, i \neq j} (\sigma^z_i -1)(\sigma^z_j - 1)\]

The first term concerns the vertices and the weights of the given pharmacores. The second term is a penalty term that penalizes edges of the graph with no interactions. The penalty \(P\) is set by the user and is defined as 6 in the cell above. The function below returns the Hamiltonina as a CUDA-Q spin_op object.

[3]:
# Generate the Hamiltonian
def ham_clique(penalty, nodes, weights, non_edges) -> cudaq.SpinOperator:

    spin_ham = 0.0
    for wt, node in zip(weights, nodes):
        #print(wt,node)
        spin_ham += 0.5 * wt * spin.z(node)
        spin_ham -= 0.5 * wt * spin.i(node)

    for non_edge in non_edges:
        u, v = (non_edge[0], non_edge[1])
        #print(u,v)
        spin_ham += penalty / 4.0 * (spin.z(u) * spin.z(v) - spin.z(u) -
                                     spin.z(v) + spin.i(u) * spin.i(v))

    return spin_ham

The code below strips the Hamiltonian into a list of coefficients and corresponding Pauli words which can be passed into a quantum kernel.

[4]:
# Collect coefficients from a spin operator so we can pass them to a kernel
def term_coefficients(ham: cudaq.SpinOperator) -> list[complex]:
    result = []
    ham.for_each_term(lambda term: result.append(term.get_coefficient()))
    return result

    # Collect Pauli words from a spin operator so we can pass them to a kernel


def term_words(ham: cudaq.SpinOperator) -> list[str]:
    result = []
    ham.for_each_term(lambda term: result.append(term.to_string(False)))
    return result


ham = ham_clique(penalty, nodes, weights, non_edges)
print(ham)

coef = term_coefficients(ham)
words = term_words(ham)

print(term_coefficients(ham))
print(term_words(ham))
[1.5+0j] IIZIIZ
[1.5+0j] ZIIZII
[-1.1657+0j] IZIIII
[1.5+0j] IZIIZI
[-1.42735+0j] IIIZII
[3.2791499999999996+0j] IIIIII
[-1.1657+0j] IIZIII
[-1.42735+0j] IIIIIZ
[-1.1657+0j] ZIIIII
[-1.42735+0j] IIIIZI

[(1.5+0j), (1.5+0j), (-1.1657+0j), (1.5+0j), (-1.42735+0j), (3.2791499999999996+0j), (-1.1657+0j), (-1.42735+0j), (-1.1657+0j), (-1.42735+0j)]
['IIZIIZ', 'ZIIZII', 'IZIIII', 'IZIIZI', 'IIIZII', 'IIIIII', 'IIZIII', 'IIIIIZ', 'ZIIIII', 'IIIIZI']

The kernel below defines a DC-QAOA circuit. What makes the approach “DC” is the inclusion of additional counteradiabatic terms to better drive the optimization to the ground state. These terms are digitized and applied as additional operations following each QAOA layer. The increase in parameters is hopefully offset by requiring fewer layers. In this example, the DC terms are additional parameterized \(Y\) operations applied to each qubit. These can be commented out to run conventional QAOA.

[5]:
@cudaq.kernel
def dc_qaoa(qubit_num:int, num_layers:int, thetas:list[float],\
    coef:list[complex], words:list[cudaq.pauli_word]):

    qubits = cudaq.qvector(qubit_num)

    h(qubits)

    count = 0
    for p in range(num_layers):

        for i in range(len(coef)):
            exp_pauli(thetas[count] * coef[i].real, qubits, words[i])
            count += 1

        for j in range(qubit_num):
            rx(thetas[count], qubits[j])
            count += 1

        #Comment out this for loop for conventional QAOA
        for k in range(qubit_num):
            ry(thetas[count], qubits[k])
            count += 1

The classical optimizer for the QAOA procedure can be specified as one of the build in CUDA-Q optimizers, in this case Nelder Mead. The parameter count is defined for DC-QAOA, but can be swapped with the commented line below for conventional QAOA.

[6]:
# Specify the optimizer and its initial parameters.
optimizer = cudaq.optimizers.NelderMead()

#Specify random seeds
np.random.seed(13)
cudaq.set_random_seed(13)

# if dc_qaoa used
parameter_count = (2 * qubit_num + len(coef)) * num_layers

# if qaoa used
# parameter_count=(qubit_num+len(coef))*num_layers

print('Total number of parameters: ', parameter_count)
optimizer.initial_parameters = np.random.uniform(-np.pi / 8, np.pi / 8,
                                                 parameter_count)
print("Initial parameters = ", optimizer.initial_parameters)
Total number of parameters:  66
Initial parameters =  [0.21810696323572243, -0.20613464375211488, 0.2546877639814583, 0.3657985647468064, 0.37118004688049144, -0.03656087558321203, 0.08564174998504231, 0.21639801853794682, 0.11122286088634259, 0.1743727097033635, -0.36518146001762486, -0.15829741539542244, -0.3467434780387345, 0.28043500852894776, -0.09986021299050934, 0.14125225086023052, -0.19141728018199775, -0.11970943368650361, -0.3853063093646483, -0.1112643868789806, 0.3527177454825464, -0.22156160012057186, -0.1418496891385843, 0.32811766468303116, -0.367642000671186, -0.34158180583996006, 0.10196745745501312, 0.29359239180502594, -0.3858537615546677, 0.19366130907065582, 0.24570488114056754, -0.3332307385378807, 0.12287973244618389, 0.007274514934614895, -0.015799547372526146, 0.3578070967202224, -0.39268963055535144, -0.19872246354138554, 0.16668715544467982, -0.13777293592446055, -0.17514665212709513, 0.15350249947988204, 0.32872977428061945, -0.20068831419712105, -0.032919322131134854, -0.19399909325771983, -0.09477141125241506, 0.08210460401106645, 0.21392577760158515, -0.3393568044538389, 0.14615087942938465, 0.03790339186006314, -0.2843250892879255, -0.3151384847055956, -0.19983741137121905, -0.27348611567665115, 0.33457528180906904, 0.14145414847455462, -0.20604220093940323, 0.05410235084309195, 0.04447870918600966, -0.3355714098595045, 0.266806440171265, -0.07436189654442632, -0.2789176729721685, -0.2427508182662484]

A cost function is specified which computes the expectation value of the DC-QAOA circuit and the Hamiltonian using the observe function. Running the optimization returns the minimized expectation value and the optimal parameters.

[7]:
cost_values = []


def objective(parameters):

    cost = cudaq.observe(dc_qaoa, ham, qubit_num, num_layers, parameters, coef,
                         words).expectation()
    cost_values.append(cost)
    return cost


# Optimize!
optimal_expectation, optimal_parameters = optimizer.optimize(
    dimensions=parameter_count, function=objective)

print('optimal_expectation =', optimal_expectation)
print('optimal_parameters =', optimal_parameters)
optimal_expectation = -2.0057493966746804
optimal_parameters = [2.0508763934174787, 0.013930789730781493, 0.5793211220774144, 0.878009560684498, 0.5277129177248182, 0.4404810513078178, 0.5755552245467919, 0.14125558672355468, 0.3724262117066903, 0.1318978057007808, -1.1228708513911436, 0.932342804955409, -0.8478237950658537, 0.46345886313018125, -0.5809397306340341, 0.2408342488137229, 0.11216088888484882, -0.009704173265255175, 0.4757346661223584, -0.7281211610985926, 0.06051951319169091, -0.7794512146826196, 0.09249435261907034, 0.09998378319110682, 1.255349350720572, 1.2607038244228248, 0.2060124032311757, 0.13991934581192997, 0.9874814082082164, -0.1591291464755939, 0.30815482837046393, -0.9701804681517978, -0.002609462845755913, 0.43533533568363353, 0.642630110681613, 0.6137063363954748, -0.7204687246344496, 0.08390768435524378, 0.5480630700433249, -0.38905723227347905, -0.6837811162838194, -0.17239016898719284, 0.1649341118754853, -0.46771209183422724, -0.008565327035838663, 1.982230359328883, -0.4232972687799105, 0.22765896988428905, 0.04207923928239914, -0.36758378917672285, -0.01825447063622079, -0.059755059728027485, -0.6849697218162497, 0.2711684382411018, -0.2904257415666667, -0.16359529445017368, -0.09168623367396612, 0.5786087806926155, -0.3476755367718726, 0.1209273564533628, 0.605136043801364, -0.19128215816141694, 0.16756583092588012, 1.0715488214105267, -0.5269641128095075, -0.3029128369198704]

Sampling the circuit with the optimal parameters allows for the most_probable command to reveal the bitsting corresponding to the ideal graph partitioning solution. This indicates what sort of interactions are present in the ideal docking configuration

[8]:
shots = 200000

counts = cudaq.sample(dc_qaoa,
                      qubit_num,
                      num_layers,
                      optimal_parameters,
                      coef,
                      words,
                      shots_count=shots)
print(counts)

print('The MVWCP is given by the partition: ', counts.most_probable())
{ 110001:16 011100:4 111000:199979 011000:1 }

The MVWCP is given by the partition:  111000

dockin

The convergence of the optimization can be plotted below.

[9]:
import matplotlib.pyplot as plt

x_values = list(range(len(cost_values)))
y_values = cost_values

plt.plot(x_values, y_values)

plt.xlabel("Epochs")
plt.ylabel("Cost Value")
plt.show()
../../_images/applications_python_digitized_counterdiabatic_qaoa_20_0.png