qmmm.ipynb

Kernel: Python 3 (ipykernel)

# SPDX-License-Identifier: Apache-2.0 AND CC-BY-NC-4.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

QM/MM Hybrid Workflows - Combining VQE with a Polarizable Embedding Framework

This lab presents a hybrid workflow called QM/MM commonly used by computational chemists to tackle challenging multi-scale problems such as chemical reaction within solution or drug-protein binding. You will learn how CUDA-Q can be used to combine quantum computing methods like VQE with classical polarizable forcefields into a single flexible workflow for QM/MM problems.

What you will do:

• Learn the concepts behind QM/MM

• Learn how to compute one-particle reduced density matrices with CUDA-Q

• Implement a QM/MM workflow using VQE for the QM part.

• Run VQE QM/MM with an active space to study a larger molecule.

• Combine ADAPT-VQE with QM/MM to improve the workflow.

CUDA-Q syntax you will use

• cudaq.spinop

• cudaq.observe

• expectation()

Prerequisites: This lab assumes learners have a basic understanding of chemical terminology and quantum computing concepts like VQE, expectation values, and qubit Hamiltonians as well as some experience working with CUDA-Q. More advanced concepts will be introduced along the way and learners will be directed to additional resources where they can dive deeper. The final exercise assumes users have completed the ADAPT-VQE CUDA-Q academic lab. If you have not done this, you can come back and finish that part later.

Run the cell below to load the packages you will be using. You may need to restart the kernel after running. Note that you may need to run apt-get update followed by apt-get install python3-dev and apt-get install build-essential to install the prerequistes needed for the CPPE library.

#may need to run the following first
#apt-get update
#apt-get install python3-dev
#apt-get install build-essential
%pip install cppe -q
%pip install pyscf -q

import sys
import os
parent_dir = os.path.join(os.getcwd(), '..')
sys.path.append(parent_dir)

from pyscf import gto, scf, mcscf, solvent
import numpy as np
import time
from functools import reduce
from scipy.optimize import minimize
import cudaq
from cudaq import spin

cudaq.set_target("nvidia", option="fp64")

from aux_files.qmmm.qchem.uccsd import uccsd_parameter_size
from aux_files.qmmm.qchem.uccsd_vqe import uccsd_circuit_vqe
from aux_files.qmmm.qchem.particle_operator import one_particle_op
from aux_files.qmmm.qchem.PEoperator import pe_operator, pe_operator_as
from aux_files.qmmm.qchem.uccsd_init_param import get_parameters
from aux_files.qmmm.qchem.hamiltonian import jordan_wigner_fermion, jordan_wigner_pe
from aux_files.qmmm.qchem.classical_pyscf import get_mol_pe_hamiltonian
from aux_files.qmmm.qchem.operator_pool import get_uccsd_pool

An overview of QM/MM

QM/MM is a multi-scale technique most useful in situations where the key properties of a chemical reaction are dominated by the interactions of a small number of atoms within an environment of many atoms. Ideally, one would simulate the entire environment with high accuracy methods. However, this requires simulation of highly complex quantum systems which quickly become too complex for even the best supercomputers.

One could choose to ignore the environment, but it does impact the results and can lead to qualitatively wrong predictions. A reasonable approximation is to split the problem into a high accuracy QM (quantum mechanics) region within a low accuracy MM (molecular mechanics/classical) region. This compromise can often result in more reliable predictions with a fraction of the compute resources.

The QM part is often simulated with something like DFT, Hartree-Fock, Coupled Cluster, etc, while the MM part is described by classical molecular mechanics using force fields. However, this depends on the specific scales in question.

Common situations for QM/MM are a drug binding to a large protein or a chemical reaction occurring in solution. The image below shows an example from "A Projector-Embedding Approach for Multiscale Coupled-Cluster Calculations Applied to Citrate Synthase" where the authors studied the first step of the citric acid cycle using high-accuracy QM simulations within the environment of an enzyme modeled with lower quality methods.

The trickiest part of QM/MM is determining how the two regions interact. All schemes introduce some sort of error. For example, one approach is called subtractive QM/MM where the entire system is computed with MM methods, then the small region alone is computed with MM methods and subtracted, finally the small region is computed with QM and added back in.

$$E_{subtractive} = E_{All}^{MM} - E_{Small}^{MM} + E_{Small}^{QM} \ne E_{All}^{QM}$$

The problem here is that both the MM and QM computations on the small region ignore interactions with the large region so error is introduced. There are a number of similar ways, but all hinge on the nature of the interaction between the regions.

A more sophisticated approach is to set up the problem such that the QM and MM regions interact through a polarizable embedding scheme. The steps for this are:

1. Define the geometry of all molecules and specify the molecules in the QM and MM regions.

2. Using an electronic structure method, compute the ground state of the QM region using an adjusted Hamiltonian that depends on the density of the MM region.

3. Use the resulting wavefunction to compute the one particle reduced density matrix (1-RDM).

4. Compute the induced dipoles in the MM region based on the 1-RDM.

5. Recompute the energy of the QM region using the new MM dipoles.

6. Repeat until self-consistent field (SCF) convergence of the energy and density.

This procedure allows the two regions to interact and provides a more realistic approximation of the physics within the QM region. To get a sense of how this works with a visualization, try the widget linked here.

Some Technical Prerequisites (Adjusted Hamiltonians and RDMs)

Construction of the QM/MM workflow requires a few technical prerequisites worth reviewing. The starting point is the standard molecular Hamiltonian for the a molecule in a vacuum ($H_0 = T_e + T_N + V_{ee} + V_{eN} + V_{NN}$), which contains the typical electron and nuclear interaction terms. However, to compute the total energy of the QM region, an additional term $H_{pe}$ is needed to capture the interaction between the electrons in the QM region and the induced dipoles of the MM region.

$$H_{QM} = H_0 + H_{pe}$$

$H_{pe}$ is the bridge which allows the SCF procedure to optimize the QM ground state and the MM region dipoles simultaneously.

Most of the technical details related to computing the induced dipoles and the $H_{pe}$ coefficients are beyond the scope for this lab and will be handled by external tools like the CPPE library. Curious readers can see the full derivation in "The variational quantum eigensolver self-consistent field method within a polarizable embedded framework".

One particular item that does warrant additional discussion is the one-particle reduced density matrix (1-RDM). This matrix contains information about all single electron properties of the QM wavefunction and is the key quantity needed from the QM wavefunction to update $H_{pe}$. Given a wavefunction $\ket{\Psi}$ composed of spin orbitals $\psi_i$, the 1-RDM is the matrix $\gamma_{pq} = \bra{\Psi} a_p^{\dagger}a_p \ket{\Psi}$.

The matrix elements can be computed using a quantum computer and evaluating the following observables with a prepared state.

If $p=q$, then $\gamma_{pp} = \frac{1}{2}(I_p - Z_p)$

If $p<q$, then $\gamma_{pq} = \frac{1}{4}(X_pX_q + Y_pY_q - iY_pX_q +i X_pY_q)\prod_{p+1}^q Z_i$

If $q<p$, then $\gamma_{pq} = \frac{1}{4}(X_qX_p + Y_qY_p + iY_qX_p - i X_qY_p)\prod_{p+1}^q Z_i$

$\gamma_{pq}$ can be decomposed into $\gamma_{pq}^{\alpha}+\gamma_{pq}^{\beta}$ where each is computed based on the index corresponding to the $\alpha$ and $\beta$ spins.

Exercise 1

Before you build the entire workflow, you need a couple of functions to compute the 1-RDM. First, define a function which returns the spin operator necessary to compute $\gamma_{pq}$. Then, define a function which takes a $\texttt{cudaq.State}$ from the preceding VQE procedure and returns the full 1-RDM in the molecular orbital basis.

Use the "TODO" flags to guide your work.

def matrix_element_operator_1_rdm(p,q):
    """
    Generates an operator for evaluating a matrix element of the 1-RDM

    Parameters:
    p (int): index of electron p
    q (int): index of electron q

    Returns:
    cudaq.SpinOperator: observable for computing gamma_pq
    """
    
    #TODO: write a function to compute the operator for computing the p,q matrix element of 1-RDM


    return qubit_op_dm

@cudaq.kernel
def optimized_psi(state: list[complex]):
    """
    Builds a CUDA-Q kernel from the state of a previous kernel

    Parameters:
    state (np.array(complex)): statevector output from cudaq.get_state

    Returns:
    CUDA-Q kernel
    """
    q = cudaq.qvector(state)


def compute_1_rdm(psi: cudaq.State, qubits_num: int, nelectrons: int):
    """
    Builds the 1-RDM from a quantum state

    Parameters:
    state (np.array(complex)): statevector output from cudaq.get_state
    qubits_num (int): number of qubits
    nelectrons (int): number of electrons

    Returns:
    np.array(float): array containing 1-RDM matrix elements.
    """

    dm_alpha = np.zeros((qubits_num // 2, qubits_num // 2))
    dm_beta = np.zeros((qubits_num // 2, qubits_num // 2))
    
    n_spatial_orbitals = qubits_num // 2
    n_occupied = nelectrons // 2
    n_virtual = n_spatial_orbitals - n_occupied   

    alpha_indices = [i * 2 for i in range(n_occupied)]
    alpha_indices += [i * 2 + nelectrons for i in range(n_virtual)]

    beta_indices = [i * 2 + 1 for i in range(n_occupied)]
    beta_indices += [i * 2 + 1 + nelectrons for i in range(n_virtual)]
        

    #TODO Compute the alpha and beta density matricies

    return dm_alpha + dm_beta
              

Building the QM/MM workflow

A minimal QM/MM example is a LiH molecule in the presence of two water molecules. You can initialize a molecule using the auxiliary get_mol_pe_hamiltonian function which takes the geometry of a molecule, and a file containing the positions of the water, their multipoles, and polarizabilities. These quantities describe the distribution of charge around the atoms and the proclivity for the electron distribution to be perturbed in an electric field.

There are two files in the cell below which only differ in the locations of the two water molecules. The cell below classically computes the Hartree Fock and CCSD energies for LiH in the presence of the solvent with the energy contribution from the solvent printed as well.

Exercise 2

Try computing the polarizable embedding energies for LiH in using the two different water potential files and comment on the magnitudes of the polarizable embedding energies. Does the result make sense?

geometry = 'Li 0.3925 0.0 0.0; H -1.1774 0.0 0.0'
water_pot= "aux_files/qmmm/qchem/4NP_in_water_far.pot"
#water_pot= "aux_files/qmmm/qchem/4NP_in_water_close.pot"


molecular_data=get_mol_pe_hamiltonian(xyz=geometry, potfile=water_pot, spin=0, charge=0, basis='sto3g', ccsd=True, verbose=True)

obi_mol = molecular_data[0]   # one body integrals
tbi_mol = molecular_data[1]   # two body integrals
e_nn = molecular_data[2]      # nuclear repulsion energy
obi_pe = molecular_data[3]    # one body integrals from PE
nelectrons = molecular_data[4] 
norbitals = molecular_data[5]

qubits_num = 2 * norbitals

#Far
# 1.0898134916455333e-05
#Close
#-0.00014871754953727954

The previous cell also saves a number of key molecular quantities obtained from the classical preprocessing. This includes the standard integrals, and the integrals from the PE model. These can be used with the function below to build a spin operator corresponding to the standard Hamiltonian (spin_mol_ham) and the additional $H_{pe}$ term (spin_pe_ham) which describes excitations caused by interactions with the water molecules.

spin_mol_ham=jordan_wigner_fermion(obi_mol, tbi_mol, e_nn, tolerance = 1e-12)
spin_pe_ham_new=jordan_wigner_pe(obi_pe, tolerance = 1e-12)

We also need the Hartree Fock coefficients from an SCF run of the molecule in the solvent. This is computed below where mf_pe produces the HF coefficients of the molecule within the polarized-embedding environment of the water molecules.

mol = gto.M(
    atom = geometry,
    spin = 0,
    charge = 0,
    basis = 'sto3g'
)

mf = scf.RHF(mol)
mf_pe = solvent.PE(mf, water_pot).run()

You are almost ready to construct the main QM/MM loop which corresponds to the figure below.

The loop will continue until the final VQE energy no longer changes between iterations, i.e. the orbitals (within the filed of the induced dipoles) is converged. One slight caveat is that the final VQE energy is not the final energy of the procedure. First, the interaction term $\bra{\theta_*}H_{pe/old}\ket{\theta_*}$ needs to be subtracted from the VQE energy, leaving only $\bra{\Psi(\theta_*)}H_{0}\ket{\Psi(\theta_*)}$. Then, the $E_{pe} = E_{MM} + \bra{\Psi(\theta_*)}H_{pe/new}\ket{\Psi(\theta_*)}$ term is added directly from the CPPE library output. This ensures the final result properly sums the energy of the QM region, the MM region, and the interaction between the two without double counting.

the cell below defines the convergence criteria for the procedure and the initial parameters for the UCCSD ansatz, obtained from the auxiliary uccsd_parameter_size function. The parameters conv_tol and vqe_tol determine when the outer SCF loop and the inner VQE procedure converges, respectively.

e_last = 0.0
conv_tol = 1e-5
dE = 1.0
cycle = 1
vqe_tol = 1e-5

singles, doubles, total = uccsd_parameter_size(nelectrons, qubits_num, spin_mult = 0)
print(f"Number of parameters: {singles} singles, {doubles} doubles, {total} total")

theta = np.zeros(total, dtype=np.float64)

Exercise 3

Now, you are ready to run the full QM/MM loop started for you below. You will be filling in the sections labeled TODO. The first step is to run VQE using the modified (PE included) Hamiltonian. The code below then saves that state which is used as an input for your `compute_1_rdm` function that returns the appropriate 1-RDM computed in the molecular orbitals basis. The following steps are needed:

1. (TODO - 1) Convert the 1-RDM from the molecular orbital (MO) basis to the atomic orbital (AO) basis. Recall this is done by performing a similarity transform where 'mf_pe' form the matrix $C$. Thus, $\gamma_{AO} = C\gamma_{MO}C^T$.

2. (Done for you) The function called $\texttt{pe\_operator}$ takes the 1-RDM in the AO basis and, using CPPE, computes a new set of induced dipoles and updates the Hamiltonian term $V_{PE}$, resulting in a new spin_pe_ham variable. It also computes a new energy of the classical MM region $E_{PE}$.

3. (TODO - 2) Compute the energy for the current step. That is $\bra{\Psi(\theta_*)}H_0 + H_{pe/old}\ket{\Psi(\theta_*)} + E_{PE/new} - \bra{\Psi(\theta_*)}H_{pe/new}\ket{\Psi(\theta_*)}$ and update the change in energy.

4. (TODO - 3) Compute the final energy after the main QM/MM loop is finished. That is, $\bra{\Psi(\theta_*)}H_0 + H_{pe/old}\ket{\Psi(\theta_*)} - \bra{\Psi(\theta_*)}H_{pe/old}\ket{\Psi(\theta_*)}E_{PE/new} + E_{pe}$.

After you complete these steps, run the code and make sure it works. Note that occasionally, VQE will get stuck trying to converge. If this happens, usually restarting fixes the problem.

while dE > conv_tol:
    spin_pe_ham = spin_pe_ham_new
    spin_ham = spin_mol_ham + spin_pe_ham

    vqe_result = uccsd_circuit_vqe(spin_mult = 0, only_singles = False, only_doubles = False, qubits_num = qubits_num,
                                   electron_count = nelectrons, optimize = True, theta = theta,
                                   hamiltonian = spin_ham, method = 'L-BFGS-B', vqe_tol = vqe_tol, verbose = False)
    
    #Save parameters to define optimized circuit
    theta = vqe_result[1]
    
    print(f"VQE Energy: {vqe_result[0]}")
    print(f"VQE converged: {vqe_result[2]}")

    if vqe_result[2]:
        final_state = np.array(vqe_result[3])  #Saves a cudaq.State from converged VQE run
        start_time = time.time()
        dm_full_mo = compute_1_rdm(final_state, qubits_num, nelectrons)
        
        end_time = time.time()
        print(end_time - start_time)

        
        #TODO 1 - convert 1-RDM from MO to AO basis.
        dm_ao =  #TODO 

        spin_pe_ham_new, E_pe, V_pe = pe_operator(dm_ao, mol, mf_pe.mo_coeff, water_pot, tolerance = 1e-12)

        #TODO 2 - Compute the energy for the current step, update the change in energy, and replace last step energy with current step
  
        print('Cycle {:d}  E_tot = {:.15g}  dE = {:g}' .format(cycle, e_last, dE), flush=True)
        print('Polarizable embedding energy: ', E_pe, flush=True)
        print('\n')
        cycle+=1
    else:
        print("Unsuccessful optimization. Restart the optimization.", flush=True)


result = uccsd_circuit_vqe(spin_mult = 0, only_singles = False, only_doubles = False, qubits_num = qubits_num,
                                   electron_count = nelectrons, optimize = False, theta = theta,
                                   hamiltonian = spin_pe_ham, method = 'L-BFGS-B', vqe_tol = vqe_tol, verbose = False)

H_pe_old = result[0]

#TODO 3 - Compute the final energy of the QM/MM procedure
final_energy = #TODO


print('\n')
print('Final result: ')
print('Polarizable embedding energy: ', E_pe)
print('PE-VQE-UCCSD energy (L-BFGS-B)= ', final_energy)

Flexibility and Scalability

One of the key benefits of QM/MM methods is flexibility. The QM region can be expanded as needed to more faithfully capture the chemistry of the system and, in principle, the QM and the MM parts of the code can swap with many different methods to optimize tradeoffs for accuracy and speed as needed.

This section will explore two such approaches for modifying your current workflow to tackle larger problems. The first demonstration will simulate the NH$_3$ molecule within solution (modeled by 46 adjacent water molecules in the MM region.). The classical cost to include the extra water molecules is minimal but can provide a better model for a solution phase molecule compared to the LiH example above with only two water molecules.

The costly change here is moving from LiH to NH$_3$ which would require 30 qubits to simulate with a 6-31G basis set. One solution is to use an active space where only a subset of the highest energy occupied orbitals and lowest energy virtual orbitals are included. This approximation greatly reduces the number of qubits while, hopefully, still capturing enough of the physics to produce a "good enough" result.

In practice, active space selection is a non-trivial task and should be undertaken with great care, especially if high accuracy is required. It can usually provide a significant reduction in cost with minimal loss in accuracy if done properly.

To visualize active space selection for the NH$_3$ molecule with a 6-31G basis set, try this widget.

Exercise 4

Look through the code below, and find the comments labeled "Look Here" and note where all of the adjustments are made to the code to account for the active space. Notice how the changes are rather minimal and do not change the overall workflow.

Try rewriting the problem with the following cases.

• 4 electrons in 4 orbitals

• 2 electrons in 4 orbitals

How much does the energy change between each scenario? Is this evidence that we are using a good approximation in this case if we are targeting chemical accuracy (0.001 hartree).

geometry = 'N  0.000   0.000   0.000; H  -0.114  -0.009   0.990; H  -0.437  -0.825  -0.358; H   0.942  -0.010  -0.222'
water_pot = 'aux_files/qmmm/qchem/46_water.pot'

#Look Here - the active space is specified when building the Hamiltonian
molecular_data = get_mol_pe_hamiltonian(xyz=geometry, potfile=water_pot, spin=0, charge=0, basis='631g',
                                          nele_cas=2, norb_cas=2, ccsd=True, verbose=True)

obi_mol = molecular_data[0]
tbi_mol = molecular_data[1]
e_core = molecular_data[2]
obi_pe = molecular_data[3]
nelectrons = molecular_data[4]
norbitals = molecular_data[5]

qubits_num = 2 * norbitals

spin_mol_ham = jordan_wigner_fermion(obi_mol, tbi_mol, e_core, tolerance = 1e-12)
spin_pe_ham_new = jordan_wigner_pe(obi_pe, tolerance = 1e-12)

mol = gto.M(
    atom = geometry,
    spin = 0,
    charge = 0,
    basis = '631g'
)

mf_pe = scf.RHF(mol)
mf_pe = solvent.PE(mf_pe, water_pot).run()

#Look Here - the coefficients for the active scape need to be obtaied classically
# to use later when building the 1-RDM
mc = mcscf.CASCI(mf_pe, norbitals, nelectrons)

e_last = 0.0
conv_tol = 1e-5
dE = 1.0
cycle = 1
vqe_tol = 1e-7

singles, doubles, total = uccsd_parameter_size(nelectrons, qubits_num, spin_mult = 0)
print(f"Number of parameters: {singles} singles, {doubles} doubles, {total} total")
theta = np.zeros(total, dtype=np.float64)

while dE > conv_tol and cycle < 20:

    spin_pe_ham = spin_pe_ham_new
    spin_ham = spin_mol_ham + spin_pe_ham

    vqe_result = uccsd_circuit_vqe(spin_mult = 0, only_singles = False, only_doubles = False, qubits_num = qubits_num,
                                     electron_count = nelectrons, optimize = True, theta = theta,
                                     hamiltonian = spin_ham, method = 'L-BFGS-B', vqe_tol = vqe_tol, verbose = False)

    theta = vqe_result[1]

    print(f"VQE Energy: {vqe_result[0]}")
    print(f"VQE converged: {vqe_result[2]}")

    if vqe_result[2]:

        final_state = np.array(vqe_result[3])

        #Look Here - The mcscf coefficients are used when transforming the orbitals from MO to AO for CPPE library 
        casdm = compute_1_rdm(final_state, qubits_num, nelectrons)

        mocore = mf_pe.mo_coeff[:, :mc.ncore]
        dm_core = np.dot(mocore, mocore.conj().T) * 2
        mocas = mf_pe.mo_coeff[:, mc.ncore:mc.ncore + mc.ncas]
        dm = dm_core + reduce(np.dot, (mocas, casdm, mocas.conj().T))

        spin_pe_ham_new, E_pe, V_pe, _ = pe_operator_as(dm, mol, mf_pe.mo_coeff, water_pot, mc)

        e_current = vqe_result[0]
        dE = np.abs(e_current - e_last)
        e_last = e_current

        print('Cycle {:d}  E_tot = {:.15g}  dE = {:g}' .format(cycle, e_last, dE), flush=True)
        print('\n')
        cycle+=1

    else:
        print("Unsuccessful optimization. Restart the optimization.", flush=True)

result = uccsd_circuit_vqe(spin_mult = 0, only_singles = False, only_doubles = False, qubits_num = qubits_num,
                             electron_count = nelectrons, optimize = False, theta = theta,
                             hamiltonian = spin_pe_ham, verbose = False)
H_pe_old = result[0]

final_energy = e_last - H_pe_old + E_pe

print('\n')
print('Final result: ')
print('Polarizable embedding energy: ', E_pe)
print('PE-VQE-UCCSD energy (L-BFGS-B)= ', final_energy)

The second example demonstrating the flexibility of QM/MM considers replacement of the VQE part of the code with a more resource efficient method like ADAPT-VQE or the Generative Quantum Eigensolver. As long as a method can produce an approximate ground state of $H_{QM}$, it can swap places with VQE with no significant impact on the rest of the QM/MM code loop, while providing potentially greater accuracy and speed.

Exercise 5

If you have already completed the ADAPT-VQE lab, try to insert your code into the QM/MM workflow. This will allow the QM/MM procedure to leverage all of the convergence benefits of ADAPT-VQE.

The code cell below is blocked where you need to input the main ADAPT loop and the ADAPT preliminary functions. Make sure you make all necessary edits to the QM/MM loop, so states and energies are handled correctly.

#------------------Setup for QM/MM--------------------
geometry = 'Li 0.3925 0.0 0.0; H -1.1774 0.0 0.0'
water_pot= "aux_files/qmmm/qchem/4NP_in_water_far.pot"
#water_pot= "aux_files/qmmm/qchem/4NP_in_water_close.pot"


molecular_data=get_mol_pe_hamiltonian(xyz=geometry, potfile=water_pot, spin=0, charge=0, basis='sto3g', ccsd=True, verbose=True)

obi_mol = molecular_data[0]   # one body integrals
tbi_mol = molecular_data[1]   # two body integrals
e_nn = molecular_data[2]      # nuclear repulsion energy
obi_pe = molecular_data[3]    # one body integrals from PE
nelectrons = molecular_data[4] 
norbitals = molecular_data[5]

qubits_num = 2 * norbitals


spin_mol_ham=jordan_wigner_fermion(obi_mol, tbi_mol, e_nn, tolerance = 1e-12)
spin_pe_ham_new=jordan_wigner_pe(obi_pe, tolerance = 1e-12)


mol = gto.M(
    atom = geometry,
    spin = 0,
    charge = 0,
    basis = 'sto3g'
)

mf = scf.RHF(mol)
mf_pe = solvent.PE(mf, water_pot).run()


e_last = 0.0
conv_tol = 1e-5
dE = 1.0
cycle = 1
vqe_tol = 1e-5

singles, doubles, total = uccsd_parameter_size(nelectrons, qubits_num, spin_mult = 0)
print(f"Number of parameters: {singles} singles, {doubles} doubles, {total} total")

theta = np.zeros(total, dtype=np.float64)

#----------------------Setup for ADAPT-VQE----------------------------------
# Functions Taken from ADAPT- VQE Lab

@cudaq.kernel
def psi(state:cudaq.State):
    q = cudaq.qvector(state)

def commutator(pools, ham):
    com_op = []
    
    for i in range(len(pools)):
        # We add the imaginary number that we excluded when generating the operator pool.
        op = 1j * pools[i]
        
        com_op.append(ham * op - op * ham)
         
    return com_op


@cudaq.kernel
def initial_state(n_qubits:int, nelectrons:int):
    qubits = cudaq.qvector(n_qubits)
    
    for i in range(nelectrons):
        x(qubits[i])



@cudaq.kernel
def kernel(theta: list[float], qubits_num: int, nelectrons: int, pool_single: list[cudaq.pauli_word], 
           coef_single: list[float], pool_double: list[cudaq.pauli_word], coef_double: list[float]):
    
    q = cudaq.qvector(qubits_num)
    
    for i in range(nelectrons):
        x(q[i])
    
    count=0
    for  i in range(0, len(coef_single), 2):
        exp_pauli(coef_single[i] * theta[count], q, pool_single[i])
        exp_pauli(coef_single[i+1] * theta[count], q, pool_single[i+1])
        count+=1

    for i in range(0, len(coef_double), 8):
        exp_pauli(coef_double[i] * theta[count], q, pool_double[i])
        exp_pauli(coef_double[i+1] * theta[count], q, pool_double[i+1])
        exp_pauli(coef_double[i+2] * theta[count], q, pool_double[i+2])
        exp_pauli(coef_double[i+3] * theta[count], q, pool_double[i+3])
        exp_pauli(coef_double[i+4] * theta[count], q, pool_double[i+4])
        exp_pauli(coef_double[i+5] * theta[count], q, pool_double[i+5])
        exp_pauli(coef_double[i+6] * theta[count], q, pool_double[i+6])
        exp_pauli(coef_double[i+7] * theta[count], q, pool_double[i+7])
        count+=1
    

#-----------------------Main QM/MM Loop-------------------------------------



import time
while dE > conv_tol:
    spin_pe_ham = spin_pe_ham_new
    spin_ham = spin_mol_ham + spin_pe_ham


#-----------------ADAPT-VQE Swapped in here---------------    

#TODO - ADD Correct Code HERE

#---------------------------------------------

    final_state = np.array(state)  
    start_time = time.time()
    dm_full_mo = compute_1_rdm(final_state, qubits_num, nelectrons)
        
    end_time = time.time()
    print(end_time - start_time)

        
    dm_ao = reduce(np.dot, (mf_pe.mo_coeff, dm_full_mo, mf_pe.mo_coeff.T))

    spin_pe_ham_new, E_pe, V_pe = pe_operator(dm_ao, mol, mf_pe.mo_coeff, water_pot, tolerance = 1e-12)

    e_current = result_vqe.fun
    dE = np.abs(e_current - e_last)
    e_last = e_current

  
    print('Cycle {:d}  E_tot = {:.15g}  dE = {:g}' .format(cycle, e_last, dE), flush=True)
    print('Polarizable embedding energy: ', E_pe, flush=True)
    print('\n')
    cycle+=1


H_pe_old = cudaq.observe(psi,spin_pe_ham, state).expectation()

final_energy = e_current - H_pe_old + E_pe

print('\n')
print('Final result: ')
print('Polarizable embedding energy: ', E_pe)
print('PE-ADAPT-VQE-UCCSD energy (L-BFGS-B)= ', final_energy)

Summary

You have just leanred how to use CUDA-Q to construct a QM/MM workflow that takes advantage of the strengths of quantum and classical computing. Advanced users are encouraged to build on this example and explore different combinations of methods based on this QM/MM workflow.