%pip install jaxopt

Requirement already satisfied: jaxopt in /home/erick.ochoalopez/Code/catalyst-latest/env/lib/python3.10/site-packages (0.7)
Requirement already satisfied: numpy>=1.18.4 in /home/erick.ochoalopez/Code/catalyst-latest/env/lib/python3.10/site-packages (from jaxopt) (1.23.5)
Requirement already satisfied: scipy>=1.0.0 in /home/erick.ochoalopez/Code/catalyst-latest/env/lib/python3.10/site-packages (from jaxopt) (1.10.0)
Requirement already satisfied: jax>=0.2.18 in /home/erick.ochoalopez/Code/catalyst-latest/env/lib/python3.10/site-packages (from jaxopt) (0.4.13)
Requirement already satisfied: jaxlib>=0.1.69 in /home/erick.ochoalopez/Code/catalyst-latest/env/lib/python3.10/site-packages (from jaxopt) (0.4.13)
Requirement already satisfied: ml-dtypes>=0.1.0 in /home/erick.ochoalopez/Code/catalyst-latest/env/lib/python3.10/site-packages (from jax>=0.2.18->jaxopt) (0.2.0)
Requirement already satisfied: opt-einsum in /home/erick.ochoalopez/Code/catalyst-latest/env/lib/python3.10/site-packages (from jax>=0.2.18->jaxopt) (3.3.0)

[notice] A new release of pip is available: 23.0.1 -> 23.2.1
[notice] To update, run: pip install --upgrade pip
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.

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)

@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

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

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:

@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[0] * 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.18415706]

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):
        diff = catalyst.grad(catalyst_cost_func, argnum=0)
        return catalyst_cost_func(params), diff(params)[0]

    opt = GradientDescent(gd_fun, stepsize=0.4, value_and_grad=True)

    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.19024097 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
def grad_descent_step(params, stepsize: float):
    diff = catalyst.grad(catalyst_cost_func, argnum=0)
    return params - diff(params)[0] * 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.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.18356271 0.18415706]

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([2], float), stepsize: float):
    diff = catalyst.grad(catalyst_cost_func, argnum=0)
    return params - diff(params)[0] * 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.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.18415706]

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)
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.012782168180745632
Excitation: [0, 1, 2, 5], Gradient: 0.0
Excitation: [0, 1, 2, 7], Gradient: 0.0
Excitation: [0, 1, 2, 9], Gradient: 0.034264503599058194
Excitation: [0, 1, 3, 4], Gradient: 0.0
Excitation: [0, 1, 3, 6], Gradient: 0.0
Excitation: [0, 1, 3, 8], Gradient: -0.03426450359905819
Excitation: [0, 1, 4, 5], Gradient: -0.023581524379110937
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.02358152437912184
Excitation: [0, 1, 6, 9], Gradient: 0.0
Excitation: [0, 1, 7, 8], Gradient: 0.0
Excitation: [0, 1, 8, 9], Gradient: -0.12362273289626591

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

# 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.005062544629162211
Excitation : [0, 4], Gradient: 0.0
Excitation : [0, 6], Gradient: 0.0
Excitation : [0, 8], Gradient: -0.0009448055879243297
Excitation : [1, 3], Gradient: 0.004926625112977063
Excitation : [1, 5], Gradient: 0.0
Excitation : [1, 7], Gradient: 0.0
Excitation : [1, 9], Gradient: 0.0014535553867885554
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
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:

@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
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.012782157554624973
Excitation: [0 1 2 5], Gradient: 0.0
Excitation: [0 1 2 7], Gradient: 0.0
Excitation: [0 1 2 9], Gradient: 0.03426451122834351
Excitation: [0 1 3 4], Gradient: 8.881784197001252e-09
Excitation: [0 1 3 6], Gradient: 2.6645352591003757e-08
Excitation: [0 1 3 8], Gradient: -0.03426446681942252
Excitation: [0 1 4 5], Gradient: -0.02358147455083781
Excitation: [0 1 4 7], Gradient: 2.6645352591003757e-08
Excitation: [0 1 4 9], Gradient: 8.881784197001252e-09
Excitation: [0 1 5 6], Gradient: 8.881784197001252e-09
Excitation: [0 1 5 8], Gradient: 8.881784197001252e-09
Excitation: [0 1 6 7], Gradient: -0.023581483432622008
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.12362274759425418

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

@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[0] * 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.04758607 -0.09855445  0.09865958  0.05330061  0.05331783  0.22938123]

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

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

doc/demos/adaptive_circuits_demo.ipynb

demos/adaptive_circuits_demo

 

Download Python script

 

Download Notebook

 

View on GitHub