[ ]:

%pip install jaxopt


# Variational Quantum Eigensolver¶

The Variational Quantum Eigensolver (VQE) is a widely used quantum algorithm with applications in quantum chemistry and portfolio optimization problems. It is an application of the Ritz variational principle, where a quantum computer is trained to prepare the ground state of a given molecule. In this demo we follow PennyLane’s demos, A brief overview of VQE and Adaptive circuits for quantum chemistry, to learn how to implement VQE and Adaptive VQE problems in Catalyst by adapting the PennyLane code.

## Importing PennyLane and Catalyst¶

In order to use PennyLane with the Catalyst compiler, we need to import several important components:

• The PennyLane framework in order to access the base QML API,

• The Catalyst Python package,

• The JAX and PennyLane versions of NumPy.

:

import catalyst
from catalyst import qjit

import pennylane as qml
from pennylane import numpy as np

import jax.numpy as jnp


## VQE for the trihydrogen cation $$H_3^{+}$$ with Catalyst¶

The algorithm takes a molecular Hamiltonian and a parametrized circuit preparing the trial state of the molecule. The cost function is defined as the expectation value of the Hamiltonian computed in the trial state. With VQE, the lowest energy state (also called the ground state) of the Hamiltonian can be computed using an iterative optimization of the cost function. In PennyLane, this optimization is performed by a classical optimizer which (in principle) leverages a quantum computer to evaluate the cost function and calculate its gradient at each optimization step.

In this section, you will learn how to implement the VQE algorithm for the trihydrogen cation $$H_3^{+}$$ (three hydrogen atoms sharing two electrons) using Catalyst. We will break the implementation into three steps:

1. Find molecular Hamiltonian for $$H_3^{+}$$.

2. Prepare trial ground step (ansatz).

3. Minimize the expectation value of the Hamiltonian.

### Standard PennyLane¶

Let’s first go over the implementation of these steps in the PennyLane demo on A brief overview of VQE.

Step 1

The first step is to specify the molecule we want to simulate. This is done by providing a list with the symbols of the constituent atoms and a one-dimensional array with the corresponding nuclear coordinates in atomic units. In the next step, we can build the electronic Hamiltonian of the hydrogen molecule using the molecular_hamiltonian() function.

:

symbols = ["H", "H", "H"]
coordinates = np.array([0.028, 0.054, 0.0, 0.986, 1.610, 0.0, 1.855, 0.002, 0.0])

# Building the molecular hamiltonian for the trihydrogen cation
hamiltonian, qubits = qml.qchem.molecular_hamiltonian(symbols, coordinates, charge=1)

print(f"qubits: {qubits}")

qubits: 6


In PennyLane, we can perform steps 2 and 3 in a few lines of code using the qml.GradientDescentOptimizer optimizer:

Step 2

:

# The Hartree-Fock State
hf = qml.qchem.hf_state(electrons=2, orbitals=6)

# Define the device, using lightning.qubit device
dev = qml.device("lightning.qubit", wires=qubits)

def cost_func(params):
qml.BasisState(hf, wires=range(qubits))
qml.DoubleExcitation(params, wires=[0, 1, 2, 3])
qml.DoubleExcitation(params, wires=[0, 1, 4, 5])
return qml.expval(hamiltonian)


Step 3

:

def workflow(params, ntrials):

for n in range(ntrials):
params, prev_energy = opt.step_and_cost(cost_func, params)
print(f"--- Step: {n}, Energy: {cost_func(params):.8f}")

return params

theta = workflow(np.array([0.0, 0.0]), 10)

print(f"Final angle parameters: {theta}")

--- Step: 0, Energy: -1.26094188
--- Step: 1, Energy: -1.26779949
--- Step: 2, Energy: -1.27110529
--- Step: 3, Energy: -1.27269203
--- Step: 4, Energy: -1.27345238
--- Step: 5, Energy: -1.27381655
--- Step: 6, Energy: -1.27399097
--- Step: 7, Energy: -1.27407452
--- Step: 8, Energy: -1.27411455
--- Step: 9, Energy: -1.27413373
Final angle parameters: [0.18356278 0.1841571 ]


### PennyLane with Catalyst¶

In Catalyst, you can define a quantum function using qml.qnode. In PennyLane, Hamiltonian objects are not yet supported as a JAX data-type, but this support is planned for the future. Hence, you need to construct the Hamiltonian using qml.Hamiltonian in your cost_func as follows:

Step 2

:

hf = qml.qchem.hf_state(electrons=2, orbitals=6)
print(f"The Hartree-Fock State: {hf}")

@qml.qnode(qml.device("lightning.qubit", wires=qubits))
def catalyst_cost_func(params):
qml.BasisState(hf, wires=range(qubits))
qml.DoubleExcitation(params, wires=[0, 1, 2, 3])
qml.DoubleExcitation(params, wires=[0, 1, 4, 5])
return qml.expval(
qml.Hamiltonian(np.array(hamiltonian.coeffs), hamiltonian.ops)
)

The Hartree-Fock State: [1 1 0 0 0 0]


Step 3

Catalyst does not provide built-in optimizers, however any optimizer that is JAX-compatible will work inside a qjit function. We will see this a bit later; here, we will implement our own simple gradient descent optimization.

We can JIT compiler the whole workflow taking advantage of the QJIT compatible for-loop:

:

@qjit
@qml.qnode(qml.device("lightning.qubit", wires=qubits))
def grad_descent(params, ntrials: int, stepsize: float):
theta = params

# for_loop can only be used in JIT mode
@catalyst.for_loop(0, ntrials, 1)
def single_step(i, theta):
h = diff(theta)
return theta - h * stepsize

return single_step(theta)

theta = grad_descent(jnp.array([0.0, 0.0]), 10, 0.4)

print(f"Final angle parameters: {theta}")

Final angle parameters: [0.18356273 0.18415705]


#### Using JAX Optimizer¶

So far, we’ve learned how to implement the gradient descent optimizer using Catalyst’s control-flows and the grad operations. However, there’s another way to calculate this optimizer using JAXopt. In the following example, we’ll utilize the jaxopt.GradientDescent function to optimize the parameters of the cost_func.” We start by importing the necessary tools from JAXopt:

:

from jaxopt import GradientDescent
from jax.lax import fori_loop


The GradientDescent method helps compute the optimizer by taking a smooth function of the form gd_fun(params, *args, **kwargs) and computing its value, or its value and its gradient, depending on the value of value_and_grad argument. To optimize params iteratively, you need to use jax.lax.fori_loop to loop over the gradient descent steps:

:

@qjit
def workflow():
def gd_fun(params):
return catalyst_cost_func(params), diff(params)

def gd_update(i, args):
(params, state) = opt.update(*args)
return (params, state)

params = jnp.array([0.0, 0.0])
state = opt.init_state(params)
upper = 10
(params, _) = fori_loop(0, upper, gd_update, (params, state))
return params

theta = workflow()

print(f"Final angle parameters: {theta}")

Final angle parameters: [0.19024096 0.19115274]


There is no significant difference between a custom implementation of the gradient descent algorithm and jaxopt.GradientDescent, since the underlying implementation is very similar.

#### JIT and AOT Compilation Modes¶

To show case just-in-time (JIT) and ahead-of-time (AOT) compilation modes in Catalyst, let’s perform a single step of the gradient descent algorithm and use the Python for-loop to optimize these parameters:

:

@qjit
return params - diff(params) * stepsize

theta = jnp.array([0.0, 0.0])

for i in range(10):
print(f"--- Step: {i}, Energy: {qjit(catalyst_cost_func)(theta):.8f}")

print(f"Final angle parameters: {theta}")

--- Step: 0, Energy: -1.26094188
--- Step: 1, Energy: -1.26779949
--- Step: 2, Energy: -1.27110529
--- Step: 3, Energy: -1.27269202
--- Step: 4, Energy: -1.27345237
--- Step: 5, Energy: -1.27381655
--- Step: 6, Energy: -1.27399097
--- Step: 7, Energy: -1.27407452
--- Step: 8, Energy: -1.27411455
--- Step: 9, Energy: -1.27413373
Final angle parameters: [0.18356273 0.18415705]


In this example, grad_descent_walk compiled at the first call (in the first iteration) and was cached effectively so that in subsequent calls you could see an order-of-magnitude speedup. Even with our custom implementation of gradient descent, this is faster than that PennyLane’s VQE implementation (with qml.GradientDescentOptimizer and PennyLane-Lightning as the backend simulator), showing another advantage of JIT compilation and Catalyst.

You can also execute the grad_descent_step in the ahead-of-time (AOT) mode to trigger compilation even before the first call by using jax.core.ShapedArray:

:

from jax.core import ShapedArray

@qjit
def grad_descent_step_aot(params: ShapedArray(, float), stepsize: float):
return params - diff(params) * stepsize

theta = jnp.array([0.0, 0.0])

for i in range(10):
print(f"--- Step: {i}, Energy: {qjit(catalyst_cost_func)(theta):.8f}")

print(f"Final angle parameters: {theta}")

--- Step: 0, Energy: -1.26094188
--- Step: 1, Energy: -1.26779949
--- Step: 2, Energy: -1.27110529
--- Step: 3, Energy: -1.27269202
--- Step: 4, Energy: -1.27345237
--- Step: 5, Energy: -1.27381655
--- Step: 6, Energy: -1.27399097
--- Step: 7, Energy: -1.27407452
--- Step: 8, Energy: -1.27411455
--- Step: 9, Energy: -1.27413373
Final angle parameters: [0.18356273 0.18415705]


## Adaptive VQE for $$LiH$$ with Catalyst¶

For certain problems, considering all possible excitations can drastically increase the cost of simulation without improving the accuracy of results. In this case, one strategy to reduce the complexity is to use adaptive circuits, customized for the molecule at hand and taking into account only those excitations that are found to be important. Applying VQE to this adapted circuit is called Adaptive VQE.

In the second part of this tutorial, you will learn how to implement the Adaptive VQE to $$LiH$$ (lithium hydride) using Catalyst.

:

symbols = ["Li", "H"]
coordinates = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 2.969280527])

# Building the molecular hamiltonian for LiH
hamiltonian, qubits = qml.qchem.molecular_hamiltonian(
symbols,
coordinates,
active_electrons=2,
active_orbitals=5
)

print(f"qubits: {qubits}")

active_electrons = 2

singles, doubles = qml.qchem.excitations(active_electrons, qubits)

print(f"single_gates: {len(singles)}")
print(f"double_gates: {len(doubles)}")

## The Hartree-Fock state
hf = qml.qchem.hf_state(active_electrons, qubits)

print(f"hf: {hf}")

qubits: 10
single_gates: 8
double_gates: 16
hf: [1 1 0 0 0 0 0 0 0 0]


In this example, there are a total of 8 single + 16 double = 24 excitations. Rather than include all of them, we will adaptively select the most relevant ones as follows:

1. Compute gradients for all double excitations.

2. Select the double excitations with gradients larger than a pre-defined threshold.

3. Perform VQE to obtain the optimized parameters for the selected double excitations.

4. Repeat steps 1 and 2 for the single excitations.

5. Perform the final VQE optimization with all the selected excitations.

### Standard PennyLane¶

Let’s first go over the implementation of these steps in the PennyLane demo on adaptive circuits for quantum chemistry:

Step 1

:

# Create a circuit that applies a selected group of gates
# to a reference Hartree-Fock state.
def circuit_1(params, excitations):
qml.BasisState(hf, wires=range(qubits))

for i, excitation in enumerate(excitations):
if len(excitation) == 4:
qml.DoubleExcitation(params[i], wires=excitation)
else:
qml.SingleExcitation(params[i], wires=excitation)
return qml.expval(hamiltonian)

# Define the device and the cost function.
dev = qml.device("lightning.qubit", wires=qubits)

# Initialize the parameter values to zero such that the gradients
# are computed w.r.t the Hartree-Fock state.
params = [0.0] * len(doubles)

for i in range(len(doubles)):

Excitation: [0, 1, 2, 3], Gradient: -0.012782168180754295
Excitation: [0, 1, 2, 5], Gradient: -6.550498368530983e-34
Excitation: [0, 1, 2, 7], Gradient: -1.3364578185469677e-34
Excitation: [0, 1, 2, 9], Gradient: -0.03426450359906927
Excitation: [0, 1, 3, 4], Gradient: -1.4482607376496288e-25
Excitation: [0, 1, 3, 6], Gradient: -7.45454212597201e-26
Excitation: [0, 1, 3, 8], Gradient: 0.03426450359906941
Excitation: [0, 1, 4, 5], Gradient: -0.023581524379116772
Excitation: [0, 1, 4, 7], Gradient: -5.145844070159646e-34
Excitation: [0, 1, 4, 9], Gradient: -8.246974456422594e-27
Excitation: [0, 1, 5, 6], Gradient: -2.256949153578792e-34
Excitation: [0, 1, 5, 8], Gradient: 9.930576275746685e-35
Excitation: [0, 1, 6, 7], Gradient: -0.023581524379135244
Excitation: [0, 1, 6, 9], Gradient: -4.3826209164064524e-27
Excitation: [0, 1, 7, 8], Gradient: 6.138901697734314e-34
Excitation: [0, 1, 8, 9], Gradient: -0.12362273289627117


Step 2

:

# Select those gates that have non-trivial gradient
doubles_select = [
doubles[i] for i in range(len(doubles)) if abs(grads[i]) > 1.0e-5
]
print(f"doubles_select: {doubles_select}")

doubles_select: [[0, 1, 2, 3], [0, 1, 2, 9], [0, 1, 3, 8], [0, 1, 4, 5], [0, 1, 6, 7], [0, 1, 8, 9]]


Step 3

:

# Add the selected gates to the circuit and perform
# one optimization step to determine the updated
# (optimized) parameters for the selected double gates.

params_doubles = np.zeros(len(doubles_select))

for n in range(20):
params_doubles = opt.step(cost_fn, params_doubles, excitations=doubles_select)

print(f"params_doubles: {params_doubles}")

params_doubles: [ 0.04758607  0.09855449 -0.09865961  0.05330068  0.05331785  0.22938117]


Step 4

:

# Keep the selected gates in the circuit and compute the gradients
# with respect to all of the single excitation gates, selecting
# those that have a non-negligible gradient.
# Repeat steps 1 and 2 for the single excitations.
def circuit_2(params, excitations, gates_select, params_select):
qml.BasisState(hf, wires=range(qubits))

for i, gate in enumerate(gates_select):
if len(gate) == 4:
qml.DoubleExcitation(params_select[i], wires=gate)
elif len(gate) == 2:
qml.SingleExcitation(params_select[i], wires=gate)

for i, gate in enumerate(excitations):
if len(gate) == 4:
qml.DoubleExcitation(params[i], wires=gate)
elif len(gate) == 2:
qml.SingleExcitation(params[i], wires=gate)
return qml.expval(hamiltonian)

params = [0.0] * len(singles)

params,
excitations=singles,
gates_select=doubles_select,
params_select=params_doubles
)

for i in range(len(singles)):

singles_select = [singles[i] for i in range(len(singles)) if abs(grads[i]) > 1.0e-5]
print(f"singles_select: {singles_select}")

Excitation : [0, 2], Gradient: 0.005062544629167019
Excitation : [0, 4], Gradient: -6.105719499217732e-18
Excitation : [0, 6], Gradient: -1.472548521003468e-17
Excitation : [0, 8], Gradient: -0.0009448055879256751
Excitation : [1, 3], Gradient: -0.004926625112981563
Excitation : [1, 5], Gradient: 4.808748891065879e-19
Excitation : [1, 7], Gradient: 9.70104692642875e-19
Excitation : [1, 9], Gradient: 0.0014535553867906221
singles_select: [[0, 2], [0, 8], [1, 3], [1, 9]]


Step 5

We now have all of the gates required to build our adaptive circuit and variationally optimize it:

:

# Perform the final VQE optimization with all the selected excitations

params = np.zeros(len(doubles_select + singles_select))

gates_select = doubles_select + singles_select

ntrials = 20
for n in range(ntrials):
params, energy = opt.step_and_cost(cost_fn, params, excitations=gates_select)

print(f"Energy: {energy:.8f}")

Energy: -7.88223735


### PennyLane with Catalyst¶

Let’s repeat these steps in Catalyst, where we can take advantage of QJIT compatible control-flow operations to perform the adaptive selection and VQE.

Step 1

Let us first create a unified circuit, merging both circuit_1 and circuit_2, before performing those 5 steps. The dynamically shaped arrays can’t be utilized in the function body of a quantum function in Catalyst, so we will need to split the list of wires for excitations into two groups of single and double excitations (double_gates and single_gates in avqe_circuit).

We then need to add a gate verifier to conditionally apply either qml.DoubleExcitation or qml.SingleExcitation, depending on the type of excitation. The first part of catalyst_avqe_circuit is used in steps 4 and 5 of our Adaptive VQE recipe, while the second part is a QJIT compatible version of circuit_1:

:

@qml.qnode(qml.device("lightning.qubit", wires=qubits))
def catalyst_avqe_circuit(
num_gates, num_selected_gates,
params, selected_params,
double_gates, single_gates,
selected_double_gates, selected_single_gates,
gates_verifier, selected_verifier,
apply_selected):
qml.BasisState(hf, wires=range(qubits))

# Used in steps 4-5
def include_selected_gates():
def apply_func(i):
param = selected_params[i]
@catalyst.cond(selected_verifier[i])
def apply_doubles():
qml.DoubleExcitation(param, selected_double_gates[i])
@apply_doubles.otherwise
def otherwise():
qml.SingleExcitation(param, selected_single_gates[i])
apply_doubles()

catalyst.for_loop(0, num_selected_gates, 1)(apply_func)()
catalyst.cond(apply_selected)(include_selected_gates)()

# Used in steps 1-5
def apply_gates(i):
@catalyst.cond(gates_verifier[i])
def apply_double():
qml.DoubleExcitation(params[i], double_gates[i])
@apply_double.otherwise
def otherwise():
qml.SingleExcitation(params[i], single_gates[i])

apply_double()
catalyst.for_loop(0, num_gates, 1)(apply_gates)()

return qml.expval(
qml.Hamiltonian(np.array(hamiltonian.coeffs), hamiltonian.ops)
)

jit_catalyst_avqe_circuit = qjit(catalyst_avqe_circuit)


If we set the initial parameter values to zero, the gradients will be computed with respect to the Hartree-Fock state:

:

@qjit
num_gates, num_selected_gates,
params, selected_params,
double_gates, single_gates,
selected_double_gates, selected_single_gates,
gates_verifier, selected_verifier,
apply_selected):

return g(num_gates, num_selected_gates,
params, selected_params,
double_gates, single_gates,
selected_double_gates, selected_single_gates,
gates_verifier, selected_verifier,
apply_selected)

num_gates = len(doubles)
params = [0.0] * num_gates
gates_verifier = [True] * num_gates
[[-1,-1,-1,-1]], [[-1,-1]], gates_verifier, [False], False)

for i in range(len(doubles)):

Excitation: [0, 1, 2, 3], Gradient: -0.012782086500351397
Excitation: [0, 1, 2, 5], Gradient: 7.993605777301127e-08
Excitation: [0, 1, 2, 7], Gradient: 7.993605777301127e-08
Excitation: [0, 1, 2, 9], Gradient: -0.03426442241050154
Excitation: [0, 1, 3, 4], Gradient: 7.993605777301127e-08
Excitation: [0, 1, 3, 6], Gradient: 1.2434497875801753e-07
Excitation: [0, 1, 3, 8], Gradient: 0.03426462669153807
Excitation: [0, 1, 4, 5], Gradient: -0.02358138573299584
Excitation: [0, 1, 4, 7], Gradient: 1.2434497875801753e-07
Excitation: [0, 1, 4, 9], Gradient: 7.993605777301127e-08
Excitation: [0, 1, 5, 6], Gradient: 7.993605777301127e-08
Excitation: [0, 1, 5, 8], Gradient: 7.993605777301127e-08
Excitation: [0, 1, 6, 7], Gradient: -0.02358138573299584
Excitation: [0, 1, 6, 9], Gradient: 3.552713678800501e-08
Excitation: [0, 1, 7, 8], Gradient: 3.552713678800501e-08
Excitation: [0, 1, 8, 9], Gradient: -0.1236226765399806


Step 2

We’ll select only those double-excitation gates with a non-trivial gradient:

:

doubles_select = [doubles[i] for i in range(len(doubles)) if abs(grads[i]) > 1.0e-5]
print(f"doubles_select: {doubles_select}")

doubles_select: [[0, 1, 2, 3], [0, 1, 2, 9], [0, 1, 3, 8], [0, 1, 4, 5], [0, 1, 6, 7], [0, 1, 8, 9]]


Step 3

Now we add the selected double gates to the circuit and perform one optimization step to determine the updated (optimized) parameters associated with them:

:

@qjit
@qml.qnode(qml.device("lightning.qubit", wires=qubits))
num_gates, num_selected_gates,
params, selected_params,
double_gates, single_gates,
selected_double_gates, selected_single_gates,
gates_verifier, selected_verifier,
apply_selected, ntrials, stepsize):

h = diff(num_gates, num_selected_gates,
theta, selected_params,
double_gates, single_gates,
selected_double_gates, selected_single_gates,
gates_verifier, selected_verifier, apply_selected)
return theta - h * stepsize

num_gates = len(doubles_select)
params_doubles = [0.0] * num_gates
gates_verifier = [True] * num_gates
ntrials = 20
stepsize = 0.5

params_doubles = grad_descent(num_gates, 1, params_doubles, [0.0],
doubles_select, singles, [[-1,-1,-1,-1]], [[-1,-1]], gates_verifier,
[False], False, ntrials, stepsize)

print(f"params_doubles: {params_doubles}")

params_doubles: [ 0.04758586  0.09855436 -0.09865972  0.05330062  0.05331779  0.22938111]


In the last two steps, we’ll utilize the selected double-excitation gates in the circuit and compute the gradients with respect to all of the single-excitation gates. After this process, we can select those single-excitation gates that have non-trivial gradients and repeat the first two steps to optimize the corresponding parameters. We then perform the final VQE optimization with all the selected single- and double-excitation gates:

Step 4

:

num_gates = len(singles)
params = [0.0] * num_gates
gates_verifier = [False] * num_gates
num_selected_gates = len(doubles_select)
selected_params = params_doubles
selected_verifier = [True] * num_selected_gates
apply_selected = True

num_gates, num_selected_gates,
params, selected_params,
doubles, singles,
doubles_select, [[0,1]],
gates_verifier, selected_verifier,
apply_selected)

singles_select = [singles[i] for i in range(len(singles)) if abs(grads[i]) > 1.0e-5]

print(f"singles_select: {singles_select}")

singles_select: [[0, 2], [0, 8], [1, 3], [1, 9]]


Step 5

All of the gates now are ready to be used in the final optimization! Here goes:

:

# Selected gates
singles_select = [[0, 2], [0, 8], [1, 3], [1, 9]]
doubles_select = [[0, 1, 2, 3], [0, 1, 2, 9], [0, 1, 3, 8], [0, 1, 4, 5], [0, 1, 6, 7], [0, 1, 8, 9]]

gates_select = doubles_select + singles_select
num_gates = len(doubles_select)
params = [0.0] * num_gates
gates_verifier = [True] * len(doubles_select) + [False] * len(singles_select)
casted_singles_select = [[-1, -1]] * len(doubles_select) + singles_select
casted_doubles_select = doubles_select
ntrials = 20
stepsize = 0.5

params = grad_descent(num_gates, 1, params, [0.0], casted_doubles_select, casted_singles_select,
[[-1,-1,-1,-1]], [[-1,-1]], gates_verifier, [False], False, ntrials, stepsize)

opt_energy = jit_catalyst_avqe_circuit(num_gates, 1, params, [0.0], casted_doubles_select,
casted_singles_select, [[-1,-1,-1,-1]], [[-1,-1]], gates_verifier, [False], False)

print(f"Energy: {opt_energy:.8f}")

Energy: -7.88202350


## Conclusion¶

In this tutorial, we’ve seen some of Catalyst’s features, including the JIT and AOT compilers and QJIT compatible gradient and control-flow operations. It’s important to remember that this compiler is still an experimental project, and there may be bugs and corner cases. That said, here are a few key takeaways:

• PennyLane compatibility: Catalyst provides support for core functionalities of PennyLane with a similar user interface. Some of the features may not be fully supported yet, but future plans include the addition of more.

• Comparable performance: The Catalyst JIT/AOT compiler has the capability to convert PennyLane programs into machine code. It is projected to yield execution performance that is comparable to that of interpreted Python code, particularly when there are multiple calls of a JIT-compiled function in a program. It is important to take into consideration that the overhead of the compilation process may impact the performance of relatively small programs.

• Reduced memory usage: Catalyst can also help to reduce the memory usage of PennyLane programs by eliminating the need to keep the original source code in memory.