```
[1]:
```

```
%pip install optax
```

```
Defaulting to user installation because normal site-packages is not writeable
Requirement already satisfied: optax in /home/tzung-han.juang/.local/lib/python3.10/site-packages (0.2.2)
Requirement already satisfied: absl-py>=0.7.1 in /home/tzung-han.juang/.local/lib/python3.10/site-packages (from optax) (2.1.0)
Requirement already satisfied: chex>=0.1.86 in /home/tzung-han.juang/.local/lib/python3.10/site-packages (from optax) (0.1.86)
Requirement already satisfied: jax>=0.1.55 in /home/tzung-han.juang/.local/lib/python3.10/site-packages (from optax) (0.4.23)
Requirement already satisfied: jaxlib>=0.1.37 in /home/tzung-han.juang/.local/lib/python3.10/site-packages (from optax) (0.4.23)
Requirement already satisfied: numpy>=1.18.0 in /home/tzung-han.juang/.local/lib/python3.10/site-packages (from optax) (1.26.4)
Requirement already satisfied: typing-extensions>=4.2.0 in /home/tzung-han.juang/.local/lib/python3.10/site-packages (from chex>=0.1.86->optax) (4.12.2)
Requirement already satisfied: toolz>=0.9.0 in /home/tzung-han.juang/.local/lib/python3.10/site-packages (from chex>=0.1.86->optax) (0.12.1)
Requirement already satisfied: ml-dtypes>=0.2.0 in /home/tzung-han.juang/.local/lib/python3.10/site-packages (from jax>=0.1.55->optax) (0.4.0)
Requirement already satisfied: opt-einsum in /home/tzung-han.juang/.local/lib/python3.10/site-packages (from jax>=0.1.55->optax) (3.3.0)
Requirement already satisfied: scipy>=1.9 in /home/tzung-han.juang/.local/lib/python3.10/site-packages (from jax>=0.1.55->optax) (1.12.0)
Note: you may need to restart the kernel to use updated packages.
```

# 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.

```
[2]:
```

```
import catalyst
from catalyst import qjit
import pennylane as qml
from pennylane import numpy as np
import jax.numpy as jnp
import warnings
warnings.filterwarnings('ignore')
```

## 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:

Find molecular Hamiltonian for \(H_3^{+}\).

Prepare trial ground step (ansatz).

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.

```
[3]:
```

```
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**

```
[4]:
```

```
# 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)
@qml.qnode(dev, diff_method="adjoint")
def cost_func(params):
qml.BasisState(hf, wires=range(qubits))
qml.DoubleExcitation(params[0], wires=[0, 1, 2, 3])
qml.DoubleExcitation(params[1], wires=[0, 1, 4, 5])
return qml.expval(hamiltonian)
```

**Step 3**

```
[5]:
```

```
def workflow(params, ntrials):
opt = qml.GradientDescentOptimizer(stepsize=0.4)
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**

```
[6]:
```

```
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[0], wires=[0, 1, 2, 3])
qml.DoubleExcitation(params[1], 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:

```
[7]:
```

```
@qjit
def grad_descent(params, ntrials: int, stepsize: float):
diff = catalyst.grad(catalyst_cost_func, argnum=0)
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.18356278 0.1841571 ]
```

#### 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 Optax. In the following example, we’ll utilize the `optax.sgd`

function to optimize the parameters of the `cost_func`

.” We start by importing the necessary tools from Optax:

```
[8]:
```

```
import optax
from jax.lax import fori_loop
```

The `sgd`

method helps compute the optimizer by taking a smooth function of the form `gd_fun(params, *args, **kwargs)`

and computing its value and its gradient. To optimize `params`

iteratively, you need to use `jax.lax.fori_loop`

to loop over the gradient descent steps:

```
[9]:
```

```
@qjit
def workflow():
def gd_fun(params):
diff = catalyst.grad(catalyst_cost_func, argnum=0)
return catalyst_cost_func(params), diff(params)
opt = optax.sgd(learning_rate=0.4)
def gd_update(i, args):
param, state = args
_, gradient = gd_fun(param)
(updates, state) = opt.update(gradient, state)
param = optax.apply_updates(param, updates)
return (param, state)
params = jnp.array([0.0, 0.0])
state = opt.init(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.18356278 0.1841571 ]
```

There is no significant difference between a custom implementation of the gradient descent algorithm and `optax.sgd`

, 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:

```
[10]:
```

```
@qjit
def grad_descent_step(params, stepsize: float):
diff = catalyst.grad(catalyst_cost_func, argnum=0)
return params - diff(params) * stepsize
theta = jnp.array([0.0, 0.0])
for i in range(10):
theta = grad_descent_step(theta, 0.4)
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.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 ]
```

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`

:

```
[11]:
```

```
from jax.core import ShapedArray
@qjit
def grad_descent_step_aot(params: ShapedArray([2], float), stepsize: float):
diff = catalyst.grad(catalyst_cost_func, argnum=0)
return params - diff(params) * stepsize
theta = jnp.array([0.0, 0.0])
for i in range(10):
theta = grad_descent_step_aot(theta, 0.4)
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.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 ]
```

## 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.

```
[12]:
```

```
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:

Compute gradients for all double excitations.

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

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

Repeat steps 1 and 2 for the single excitations.

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**

```
[13]:
```

```
# 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)
cost_fn = qml.QNode(circuit_1, dev, diff_method="adjoint")
circuit_gradient = qml.grad(cost_fn, argnum=0)
# Initialize the parameter values to zero such that the gradients
# are computed w.r.t the Hartree-Fock state.
params = [0.0] * len(doubles)
grads = circuit_gradient(params, excitations=doubles)
for i in range(len(doubles)):
print(f"Excitation: {doubles[i]}, Gradient: {grads[i]}")
```

```
Excitation: [0, 1, 2, 3], Gradient: -0.012782168180745743
Excitation: [0, 1, 2, 5], Gradient: 0.0
Excitation: [0, 1, 2, 7], Gradient: 0.0
Excitation: [0, 1, 2, 9], Gradient: 0.03426450359905824
Excitation: [0, 1, 3, 4], Gradient: 0.0
Excitation: [0, 1, 3, 6], Gradient: 0.0
Excitation: [0, 1, 3, 8], Gradient: -0.034264503599058235
Excitation: [0, 1, 4, 5], Gradient: -0.023581524379111044
Excitation: [0, 1, 4, 7], Gradient: 0.0
Excitation: [0, 1, 4, 9], Gradient: 0.0
Excitation: [0, 1, 5, 6], Gradient: 0.0
Excitation: [0, 1, 5, 8], Gradient: 0.0
Excitation: [0, 1, 6, 7], Gradient: -0.02358152437912195
Excitation: [0, 1, 6, 9], Gradient: 0.0
Excitation: [0, 1, 7, 8], Gradient: 0.0
Excitation: [0, 1, 8, 9], Gradient: -0.12362273289626559
```

**Step 2**

```
[14]:
```

```
# 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**

```
[15]:
```

```
# Add the selected gates to the circuit and perform
# one optimization step to determine the updated
# (optimized) parameters for the selected double gates.
opt = qml.GradientDescentOptimizer(stepsize=0.5)
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**

```
[16]:
```

```
# 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)
cost_fn = qml.QNode(circuit_2, dev, diff_method="adjoint")
circuit_gradient = qml.grad(cost_fn, argnum=0)
params = [0.0] * len(singles)
grads = circuit_gradient(
params,
excitations=singles,
gates_select=doubles_select,
params_select=params_doubles
)
for i in range(len(singles)):
print(f"Excitation : {singles[i]}, Gradient: {grads[i]}")
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.005062544629162063
Excitation : [0, 4], Gradient: 0.0
Excitation : [0, 6], Gradient: 0.0
Excitation : [0, 8], Gradient: -0.0009448055879247612
Excitation : [1, 3], Gradient: 0.004926625112976861
Excitation : [1, 5], Gradient: 0.0
Excitation : [1, 7], Gradient: 0.0
Excitation : [1, 9], Gradient: 0.0014535553867889794
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:

```
[17]:
```

```
# Perform the final VQE optimization with all the selected excitations
cost_fn = qml.QNode(circuit_1, dev, diff_method="adjoint")
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`

:

```
[18]:
```

```
@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:

```
[19]:
```

```
@qjit
def avqe_circuit_grad(
num_gates, num_selected_gates,
params, selected_params,
double_gates, single_gates,
selected_double_gates, selected_single_gates,
gates_verifier, selected_verifier,
apply_selected):
g = catalyst.grad(catalyst_avqe_circuit, argnum=2, method="fd")
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 = jnp.asarray([0.0] * num_gates)
gates_verifier = jnp.asarray([True] * num_gates)
doubles = jnp.asarray(doubles)
singles = jnp.asarray(singles)
gates_verifier = jnp.asarray(gates_verifier)
grads = avqe_circuit_grad(num_gates, 1, params, jnp.asarray([0.0]), doubles, singles,
jnp.asarray([[-1,-1,-1,-1]]), jnp.asarray([[-1,-1]]), gates_verifier, jnp.asarray([False]), False)
for i in range(len(doubles)):
print(f"Excitation: {doubles[i]}, Gradient: {grads[i]}")
```

```
Excitation: [0 1 2 3], Gradient: -0.012782148672840776
Excitation: [0 1 2 5], Gradient: 1.7763568394002505e-08
Excitation: [0 1 2 7], Gradient: 1.7763568394002505e-08
Excitation: [0 1 2 9], Gradient: 0.034264520110127705
Excitation: [0 1 3 4], Gradient: 1.7763568394002505e-08
Excitation: [0 1 3 6], Gradient: 1.7763568394002505e-08
Excitation: [0 1 3 8], Gradient: -0.03426448458299092
Excitation: [0 1 4 5], Gradient: -0.023581492314406205
Excitation: [0 1 4 7], Gradient: 1.7763568394002505e-08
Excitation: [0 1 4 9], Gradient: 1.7763568394002505e-08
Excitation: [0 1 5 6], Gradient: 1.7763568394002505e-08
Excitation: [0 1 5 8], Gradient: 1.7763568394002505e-08
Excitation: [0 1 6 7], Gradient: -0.023581501196190402
Excitation: [0 1 6 9], Gradient: 8.881784197001252e-09
Excitation: [0 1 7 8], Gradient: 8.881784197001252e-09
Excitation: [0 1 8 9], Gradient: -0.1236227031853332
```

**Step 2**

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

```
[20]:
```

```
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: [Array([0, 1, 2, 3], dtype=int64), Array([0, 1, 2, 9], dtype=int64), Array([0, 1, 3, 8], dtype=int64), Array([0, 1, 4, 5], dtype=int64), Array([0, 1, 6, 7], dtype=int64), Array([0, 1, 8, 9], dtype=int64)]
```

**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:

```
[21]:
```

```
@qjit
def grad_descent(
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):
diff = catalyst.grad(catalyst_avqe_circuit, argnum=2, method="fd")
def grad_step(i, theta):
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
return catalyst.for_loop(0, ntrials, 1)(grad_step)(params)
num_gates = len(doubles_select)
params_doubles = jnp.asarray([0.0] * num_gates)
gates_verifier = jnp.asarray([True] * num_gates)
doubles_select = jnp.asarray(doubles_select)
singles = jnp.asarray(singles)
ntrials = 20
stepsize = 0.5
params_doubles = grad_descent(num_gates, 1, params_doubles, jnp.asarray([0.0]),
doubles_select, singles, jnp.asarray([[-1,-1,-1,-1]]), jnp.asarray([[-1,-1]]), gates_verifier,
jnp.asarray([False]), False, ntrials, stepsize)
print(f"params_doubles: {params_doubles}")
```

```
params_doubles: [ 0.04758599 -0.09855459 0.09865966 0.05330077 0.05331789 0.22938115]
```

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**

```
[22]:
```

```
num_gates = len(singles)
params = jnp.asarray([0.0] * num_gates)
gates_verifier = jnp.asarray([False] * num_gates)
num_selected_gates = len(doubles_select)
selected_params = jnp.asarray(params_doubles)
selected_verifier = jnp.asarray([True] * num_selected_gates)
apply_selected = True
grads = avqe_circuit_grad(
num_gates, num_selected_gates,
params, selected_params,
doubles, singles,
doubles_select, jnp.asarray([[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: [Array([0, 2], dtype=int64), Array([0, 8], dtype=int64), Array([1, 3], dtype=int64), Array([1, 9], dtype=int64)]
```

**Step 5**

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

```
[23]:
```

```
# 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
#print(gates_select)
num_gates = len(doubles_select)
params = jnp.asarray([0.0] * num_gates)
gates_verifier = jnp.asarray([True] * len(doubles_select) + [False] * len(singles_select))
casted_singles_select = jnp.asarray([[-1, -1]] * len(doubles_select) + singles_select)
casted_doubles_select = jnp.asarray(doubles_select)
ntrials = 20
stepsize = 0.5
params = grad_descent(num_gates, 1, params, jnp.asarray([0.0]), casted_doubles_select, casted_singles_select,
jnp.asarray([[-1,-1,-1,-1]]), jnp.asarray([[-1,-1]]), gates_verifier, jnp.asarray([False]), False, ntrials, stepsize)
opt_energy = jit_catalyst_avqe_circuit(num_gates, 1, params, jnp.asarray([0.0]), casted_doubles_select,
casted_singles_select, jnp.asarray([[-1,-1,-1,-1]]), jnp.asarray([[-1,-1]]), gates_verifier, jnp.asarray([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.