Variational quantum linear solver

Hello! If applicable, put your complete code example down below. Make sure that your code:

  • is 100% self-contained — someone can copy-paste exactly what is here and run it to
    reproduce the behaviour you are observing
  • includes comments
# Put code here
# -*- coding: utf-8 -*-

# -*- coding: utf-8 -*-
"""Improved VQLS Implementation"""

# Pennylane
import pennylane as qml
from pennylane import numpy as np

# Plotting
import matplotlib.pyplot as plt
import seaborn as sns

# Set random seed for reproducibility
np.random.seed(0)

"""## Hyperparameters"""
n_qubits = 3  # Number of system qubits
n_shots = 10 ** 6  # Number of quantum measurements
tot_qubits = n_qubits + 1  # Addition of an ancillary qubit
ancilla_idx = n_qubits  # Index of the ancillary qubit (last position)
steps = 60  # Number of optimization steps
eta = 5e-3  # Adjusted learning rate (originally 0.8)
q_delta = 0.001  # Initial spread of random quantum weights
n_layers = 1  # Number of layers in the variational circuit

"""## Problem Definition"""
# Define the matrix A (8x8 example)
A = np.array([
    [1. , 0.6 , 0. , 0. , 0., 0. , 0. , 0. ],
    [0.6 , 1. , 0.6 , 0. , 0. , 0., 0. , 0. ],
    [0. , 0.6 , 1. , 0. , 0. , 0. , 0. , 0. ],
    [0. , 0. , 0. , 1. , 0. , 0. , 0. , 0. ],
    [0., 0. , 0. , 0. , 1. , 0. , 0. , 0. ],
    [0. , 0., 0. , 0. , 0. , 1. , 0. , 0. ],
    [0. , 0. , 0. , 0. , 0. , 0. , 1. , 0. ],
    [0. , 0. , 0. , 0. , 0. , 0. , 0. , 1. ]
])

# Decompose A into Pauli terms
LCU = qml.pauli_decompose(A)
LCU_coeffs, LCU_ops = LCU.terms()
c = np.array(LCU_coeffs)

"""## Quantum Circuits"""

def U_b():
    """Unitary matrix rotating the ground state to the problem vector |b> = U_b |0>."""
    for idx in range(n_qubits):
        qml.Hadamard(wires=idx)

def CA(idx):
    """Controlled versions of the unitary components A_l of the problem matrix A."""
    if idx == 0:
        None  # Identity operation
    elif idx == 1:
        qml.CNOT(wires=[ancilla_idx, 2])
    elif idx == 2:
        qml.CNOT(wires=[ancilla_idx, 1])
        qml.CNOT(wires=[ancilla_idx, 2])
    elif idx == 3:
        qml.CY(wires=[ancilla_idx, 1])
        qml.CY(wires=[ancilla_idx, 2])
    elif idx == 4:
        qml.CZ(wires=[ancilla_idx, 1])
        qml.CNOT(wires=[ancilla_idx, 2])
    elif idx == 5:
        qml.CZ(wires=[ancilla_idx, 0])
        qml.CNOT(wires=[ancilla_idx, 2])
    elif idx == 6:
        qml.CZ(wires=[ancilla_idx, 0])
        qml.CNOT(wires=[ancilla_idx, 1])
        qml.CNOT(wires=[ancilla_idx, 2])
    elif idx == 7:
        qml.CZ(wires=[ancilla_idx, 0])
        qml.CY(wires=[ancilla_idx, 1])
        qml.CY(wires=[ancilla_idx, 2])
    elif idx == 8:
        qml.CZ(wires=[ancilla_idx, 0])
        qml.CZ(wires=[ancilla_idx, 1])
        qml.CNOT(wires=[ancilla_idx, 2])


"""## Enhanced Variational Circuit"""

# def variational_block(weights):
#     """
#     Custom variational circuit with RY rotations and CNOT entanglers.
#     Args:
#         weights: Array of shape (2*n_qubits,) for rotation parameters
#     """
#     # Apply RY rotations with the first set of parameters
#     for i in range(n_qubits):
#         qml.RY(weights[i], wires=i)

#     # Apply CNOT gates with adjacent qubits (cyclically connected)
#     for i in range(n_qubits):
#         qml.CNOT(wires=[i, (i + 1) % n_qubits])

    # # Apply RY rotations with the second set of parameters
    # for i in range(n_qubits, 2*n_qubits):
    #     qml.RY(weights[i], wires=i % n_qubits)

    # # Apply CNOT gates in reverse order
    # for i in reversed(range(n_qubits)):
    #     qml.CNOT(wires=[i, (i + 1) % n_qubits])

def variational_block(weights):
    """
    Enhanced variational circuit with StronglyEntanglingLayers.
    Args:
        weights: Array of shape (n_layers, n_qubits, 3) for rotation parameters
    """
    # Initial Hadamard layer for superposition
    for idx in range(n_qubits):
        qml.Hadamard(wires=idx)

    # Strongly entangling layers for variational optimization
    qml.StronglyEntanglingLayers(weights, wires=range(n_qubits))


"""## Hadamard Test Implementation"""

dev_mu = qml.device("lightning.qubit", wires=tot_qubits)

@qml.qnode(dev_mu, interface="autograd")
def local_hadamard_test(weights, l=None, lp=None, j=None, part=None):
    """Hadamard test for estimating expectation values."""
    # First Hadamard gate on ancillary qubit
    qml.Hadamard(wires=ancilla_idx)

    # For imaginary part, add a -i phase gate
    if part == "Im" or part == "im":
        qml.PhaseShift(-np.pi / 2, wires=ancilla_idx)

    # Variational circuit
    variational_block(weights)

    # Controlled unitary A_l
    CA(l)

    # Adjoint of U_b (which is U_b itself in this case)
    U_b()

    # Controlled Z operator if j != -1
    if j != -1:
        qml.CZ(wires=[ancilla_idx, j])

    # Unitary U_b
    U_b()

    # Controlled adjoint(A_lp)
    CA(lp)

    # Final Hadamard on ancillary
    qml.Hadamard(wires=ancilla_idx)

    # Expectation value of Z on ancillary
    return qml.expval(qml.PauliZ(wires=ancilla_idx))

def mu(weights, l=None, lp=None, j=None):
    """Compute complex coefficients for cost function."""
    mu_real = local_hadamard_test(weights, l=l, lp=lp, j=j, part="Re")
    mu_imag = local_hadamard_test(weights, l=l, lp=lp, j=j, part="Im")
    return mu_real + 1.0j * mu_imag

"""## Cost Functions"""

def psi_norm(weights):
    """Compute <psi|psi> where |psi> = A|x>."""
    norm = 0.0
    for l in range(len(c)):
        for lp in range(len(c)):
            norm += c[l] * np.conj(c[lp]) * mu(weights, l, lp, -1)
    return abs(norm)

def cost_loc(weights):
    """Local cost function that tends to zero when A|x> ∝ |b>."""
    mu_sum = 0.0
    for l in range(len(c)):
        for lp in range(len(c)):
            for j in range(n_qubits):
                mu_sum += c[l] * np.conj(c[lp]) * mu(weights, l, lp, j)
    mu_sum = abs(mu_sum)
    return 0.5 - 0.5 * mu_sum / (n_qubits * psi_norm(weights))

"""## Optimization"""

# Initialize weights with proper shape for our custom variational circuit
w = q_delta * np.random.randn(n_layers, n_qubits, 3, requires_grad=True)

# Use Adam optimizer for better convergence
opt = qml.AdamOptimizer(eta)

cost_history = []
for it in range(steps):
    w, cost = opt.step_and_cost(cost_loc, w)
    print(f"Step {it:3d} Cost_L = {cost:9.7f}")
    cost_history.append(cost)

# Plotting
plt.style.use("seaborn-v0_8")
plt.plot(cost_history, "g")
plt.ylabel("Cost function")
plt.xlabel("Optimization steps")
plt.title(f"Learning rate η = {eta}")
plt.show()

"""## Classical Solution for Comparison"""

# Pauli matrices
Id = np.identity(2)
Z = np.array([[1, 0], [0, -1]])
X = np.array([[0, 1], [1, 0]])
Y = np.array([[0, -1j], [1j, 0]])

# Construct A from Pauli terms
A_0 = np.identity(8)
A_1 = np.kron(np.kron(Id, Id), X)
A_2 = np.kron(np.kron(Id, X), X)
A_3 = np.kron(np.kron(Id, Y), Y)
A_4 = np.kron(np.kron(Id, Z), X)
A_5 = np.kron(np.kron(Z, Id), X)
A_6 = np.kron(np.kron(Z, X), X)
A_7 = np.kron(np.kron(Z, Y), Y)
A_8 = np.kron(np.kron(Z, Z), X)

A_num = sum(c[i] * A_i for i, A_i in enumerate([A_0, A_1, A_2, A_3, A_4, A_5, A_6, A_7, A_8]))
b = np.ones(8) / np.sqrt(8)  # |b> state

# Classical solution
A_inv = np.linalg.inv(A_num)
x_classical = np.dot(A_inv, b)
c_probs = (x_classical / np.linalg.norm(x_classical)) ** 2

"""## Quantum Solution"""

dev_x = qml.device("lightning.qubit", wires=n_qubits, shots=n_shots)

@qml.qnode(dev_x, interface="autograd")
def prepare_and_sample(weights):
    """Prepare the quantum state and sample from it."""
    variational_block(weights)
    return qml.sample()

raw_samples = prepare_and_sample(w)
samples = [int("".join(str(bs) for bs in sam), base=2) for sam in raw_samples]
q_probs = np.bincount(samples, minlength=2**n_qubits) / n_shots

"""## Results Comparison"""

print("Classical solution probabilities:\n", c_probs)
print("\nQuantum solution probabilities:\n", q_probs)

# Plotting
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))

ax1.bar(np.arange(2**n_qubits), c_probs, color="blue")
ax1.set_xlim(-0.5, 2**n_qubits - 0.5)
ax1.set_xlabel("Basis state")
ax1.set_ylabel("Probability")
ax1.set_title("Classical Solution")

ax2.bar(np.arange(2**n_qubits), q_probs, color="green")
ax2.set_xlim(-0.5, 2**n_qubits - 0.5)
ax2.set_xlabel("Basis state")
ax2.set_ylabel("Probability")
ax2.set_title("Quantum Solution")

plt.tight_layout()
plt.show()

If you want help with diagnosing an error, please put the full error message below:

Why is the classical solution different from the quantum one?
# Put full error message here
Why is the classical solution different from the quantum one?
And, finally, make sure to include the versions of your packages. Specifically, show us the output of `qml.about()`.

Hi @moslem_ym
Welcome to the forum!
I see you have adapted the code from this tutorial on VQLS.
I was able to run your code and see what you mean by the two solutions not matching each other.

Since in the tutorial example the two solutions do match, could you please provide details on what the differences between your code and the original one are? This will speed up things and help me to help you.

Thank you.

Hi @daniela.murcillo

I tried out the second method, and the matrix from the tutorial is sparse. However, it seems that this code isn’t quite achieving exact convergence for the three-diagonal matrix I used, and the results are quite different from the classic approach.
Thank you

Hi!
Yes, now I see. Your system of linear equations is more complicated and for that reason, you used StronglyEntanglingLayers for your VQC. Your code works nicely and I think you implemented this the right way.
This is a tricky question because being a variational algorithm that leaves so much open, it is not guaranteed that it will converge.
However, I think that adjusting some of the parameters can help. I did some trial an error and manage to get lower values of the cost (and better convergence) and the solution looks a bit better now. You might be able to improve this further and get better results. The cost started at 0.044 and ended at 0.0042.
I used more steps, two layers and a bigger delta.

steps = 100  # Number of optimization steps
eta = 5e-3  # Adjusted learning rate (originally 0.8)
q_delta = 0.1  # Initial spread of random quantum weights
n_layers = 2  # Number of layers in the variational circuit

If you are curious about this subject there is also another recent demo for solving linear system problems.

I hope this was helpful.