Simulate the variational quantum eigensolver (VQE)

In this guide, we show how to use ffsim to simulate the variational quantum eigensolver (VQE). Here we use the local unitary cluster Jastrow (LUCJ) ansatz to produce the variationally optimized quantum state, but the workflow is similar for other variational ansatzes in ffsim, and you can also define your own ansatz. We’ll use VQE to calculate an approximation to the ground state energy of a nitrogen molecule in an active space of 6 electrons in 6 orbitals derived from the STO-6G basis set.

First, let’s build the molecule.

import pyscf
import pyscf.mcscf

import ffsim

# Build N2 molecule
mol = pyscf.gto.Mole()
mol.build(
    atom=[["N", (0, 0, 0)], ["N", (1.0, 0, 0)]],
    basis="sto-6g",
    symmetry="Dooh",
)

# Define active space
n_frozen = 4
active_space = range(n_frozen, mol.nao_nr())

# Get molecular data and molecular Hamiltonian (one- and two-body tensors)
scf = pyscf.scf.RHF(mol).run()
mol_data = ffsim.MolecularData.from_scf(scf, active_space=active_space)
norb = mol_data.norb
nelec = mol_data.nelec
mol_hamiltonian = mol_data.hamiltonian

# Compute FCI energy
cas = pyscf.mcscf.CASCI(scf, ncas=norb, nelecas=nelec).run()

print(f"FCI energy = {cas.e_tot}")
print(f"norb = {norb}")
print(f"nelec = {nelec}")

WARN: Unable to to identify input symmetry using original axes.
Different symmetry axes will be used.

converged SCF energy = -108.464957764796
CASCI E = -108.566842251942  E(CI) = -11.9110176528507  S^2 = 0.0000000
FCI energy = -108.56684225194182
norb = 6
nelec = (3, 3)

General UCJ ansatz

Since our molecule has a closed-shell Hartree-Fock state, we’ll use the spin-balanced variant of the UCJ ansatz, UCJOpSpinBalanced. We’ll initialize the ansatz from t2 amplitudes obtained from a CCSD calculations. We’ll first demonstrate the general UCJ ansatz, without adding the locality constraints of the LUCJ ansatz just yet.

The following code cell initializes the ansatz operator, applies it to the Hartree-Fock state, and computes the energy of the resulting state.

import numpy as np
from pyscf import cc

# Get CCSD t2 amplitudes for initializing the ansatz
ccsd = cc.CCSD(
    scf, frozen=[i for i in range(mol.nao_nr()) if i not in active_space]
).run()

# Construct UCJ operator
n_reps = 2
operator = ffsim.UCJOpSpinBalanced.from_t_amplitudes(ccsd.t2, t1=ccsd.t1, n_reps=n_reps)

# Construct the Hartree-Fock state to use as the reference state
reference_state = ffsim.hartree_fock_state(norb, nelec)

# Apply the operator to the reference state
ansatz_state = ffsim.apply_unitary(reference_state, operator, norb=norb, nelec=nelec)

# Compute the energy ⟨ψ|H|ψ⟩ of the ansatz state
hamiltonian = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)
energy = np.real(np.vdot(ansatz_state, hamiltonian @ ansatz_state))
print(f"Energy at initialization: {energy}")

E(CCSD) = -108.5658290955831  E_corr = -0.1008713307875627
Energy at initialization: -108.55281964833208

To variationally optimize the ansatz, we’ll take advantage of methods for conversion to and from real-valued parameter vectors. In the following code cell, we define an objective function that takes a parameter vector as input and outputs the energy of the associated ansatz state. We then optimize this objective function using scipy.optimize.minimize, with an initial guess obtained from the operator we initialized previously from t2 amplitudes.

import scipy.optimize


def fun(x):
    # Initialize the ansatz operator from the parameter vector
    operator = ffsim.UCJOpSpinBalanced.from_parameters(
        x, norb=norb, n_reps=n_reps, with_final_orbital_rotation=True
    )
    # Apply the ansatz operator to the reference state
    final_state = ffsim.apply_unitary(reference_state, operator, norb=norb, nelec=nelec)
    # Return the energy ⟨ψ|H|ψ⟩ of the ansatz state
    return np.real(np.vdot(final_state, hamiltonian @ final_state))


result = scipy.optimize.minimize(
    fun, x0=operator.to_parameters(), method="L-BFGS-B", options=dict(maxiter=10)
)

print(f"Number of parameters: {len(result.x)}")
print(result)

Number of parameters: 192
  message: STOP: TOTAL NO. OF ITERATIONS REACHED LIMIT
  success: False
   status: 1
      fun: -108.56163391391395
        x: [ 1.027e-01  1.308e+00 ... -1.282e-03 -1.337e-02]
      nit: 10
      jac: [ 4.405e-05 -3.979e-05 ... -1.748e-04 -8.640e-04]
     nfev: 2509
     njev: 13
 hess_inv: <192x192 LbfgsInvHessProduct with dtype=float64>

LUCJ ansatz

Now, let’s add locality constraints to simulate the LUCJ ansatz. We’ll restrict same-spin interactions to a line topology, and opposite-spin interactions to those within the same spatial orbital. As explained in The local unitary cluster Jastrow (LUCJ) ansatz, these constraints allow the ansatz to be simulated directly on a square lattice.

In the following code cell, we demonstrate the optimization of the ansatz with these constraints imposed. Notice that with the same number of ansatz repetitions, the number of parameters of the ansatz has decreased from 88 to 62.

pairs_aa = [(p, p + 1) for p in range(norb - 1)]
pairs_ab = [(p, p) for p in range(norb)]
interaction_pairs = (pairs_aa, pairs_ab)


def fun(x):
    operator = ffsim.UCJOpSpinBalanced.from_parameters(
        x,
        norb=norb,
        n_reps=n_reps,
        interaction_pairs=interaction_pairs,
        with_final_orbital_rotation=True,
    )
    final_state = ffsim.apply_unitary(reference_state, operator, norb=norb, nelec=nelec)
    return np.real(np.vdot(final_state, hamiltonian @ final_state))


result = scipy.optimize.minimize(
    fun,
    x0=operator.to_parameters(interaction_pairs=interaction_pairs),
    method="L-BFGS-B",
    options=dict(maxiter=10),
)
print(f"Number of parameters: {len(result.x)}")
print(result)

Number of parameters: 130
  message: STOP: TOTAL NO. OF ITERATIONS REACHED LIMIT
  success: False
   status: 1
      fun: -108.53960328640134
        x: [ 1.059e-01  1.310e+00 ... -5.941e-04 -1.075e-02]
      nit: 10
      jac: [-3.084e-04 -4.675e-04 ... -2.998e-04 -1.151e-04]
     nfev: 1703
     njev: 13
 hess_inv: <130x130 LbfgsInvHessProduct with dtype=float64>

Optimize with the linear method

ffsim includes an implementation of the “linear method” for optimization of a variational wavefunction. The linear method often converges faster than a standard optimization algorithm like L-BFGS-B. The interface is similar to that of scipy.optimize.minimize, the main difference being that instead of passing a callable that directly returns the function value to be optimized, you pass two objects: a callable that returns the wavefunction, and the Hamiltonian representing the energy to be optimized as a LinearOperator. The code cell below shows how to use the linear method to optimize the LUCJ ansatz from the previous example. It also shows how you can use an optional callback function to save intermediate results of the optimization.

from collections import defaultdict

from ffsim.optimize import minimize_linear_method


# Define function that converts a list of parameters to the corresponding state vector
def params_to_vec(x: np.ndarray) -> np.ndarray:
    operator = ffsim.UCJOpSpinBalanced.from_parameters(
        x,
        norb=norb,
        n_reps=n_reps,
        interaction_pairs=interaction_pairs,
        with_final_orbital_rotation=True,
    )
    return ffsim.apply_unitary(reference_state, operator, norb=norb, nelec=nelec)


# Define a callback function used to save optimization information (this is optional)
info = defaultdict(list)


def callback(intermediate_result: scipy.optimize.OptimizeResult):
    # The callback function is called after each iteration. It accepts
    # an OptimizeResult object storing the parameters and function value at
    # the current iteration, and possibly other information
    info["x"].append(intermediate_result.x)
    info["fun"].append(intermediate_result.fun)
    if hasattr(intermediate_result, "jac"):
        info["jac"].append(intermediate_result.jac)
    if hasattr(intermediate_result, "regularization"):
        info["regularization"].append(intermediate_result.regularization)
    if hasattr(intermediate_result, "variation"):
        info["variation"].append(intermediate_result.variation)


# Optimize with the linear method
result = minimize_linear_method(
    params_to_vec,
    hamiltonian,
    x0=operator.to_parameters(interaction_pairs=interaction_pairs),
    maxiter=10,
    callback=callback,
)

# Print some information
print(f"Number of parameters: {len(result.x)}")
print(result)
print()
for i, (fun, jac, regularization, variation) in enumerate(
    zip(info["fun"], info["jac"], info["regularization"], info["variation"])
):
    print(f"Iteration {i + 1}")
    print(f"    Energy: {fun}")
    print(f"    Norm of gradient: {np.linalg.norm(jac)}")
    print(f"    Regularization hyperparameter: {np.linalg.norm(regularization)}")
    print(f"    Variation hyperparameter: {np.linalg.norm(variation)}")

Number of parameters: 130
 message: Stop: Total number of iterations reached limit.
 success: False
     fun: -108.56076744181719
       x: [ 1.355e-02  1.128e+00 ...  6.895e-02 -3.678e-01]
     nit: 10
     jac: [-9.210e-04 -2.212e-04 ... -1.019e-03 -5.295e-04]
    nfev: 3025
    njev: 10
  nlinop: 1725

Iteration 1
    Energy: -108.53954720324322
    Norm of gradient: 0.013769241695895162
    Regularization hyperparameter: 0.005511999506582103
    Variation hyperparameter: 0.9896282148418554
Iteration 2
    Energy: -108.53991280839396
    Norm of gradient: 0.013145989057117039
    Regularization hyperparameter: 0.002253910015290631
    Variation hyperparameter: 0.9896535534840212
Iteration 3
    Energy: -108.54196380907202
    Norm of gradient: 0.04980184611578818
    Regularization hyperparameter: 0.0028530110411914924
    Variation hyperparameter: 0.989653385439193
Iteration 4
    Energy: -108.54790605237028
    Norm of gradient: 0.035607734899118025
    Regularization hyperparameter: 0.009697879303123847
    Variation hyperparameter: 0.989652885177815
Iteration 5
    Energy: -108.552425747902
    Norm of gradient: 0.07113327576212645
    Regularization hyperparameter: 0.004626931177017402
    Variation hyperparameter: 0.9896516962839016
Iteration 6
    Energy: -108.55681501112797
    Norm of gradient: 0.07599847484727296
    Regularization hyperparameter: 0.004086292816093098
    Variation hyperparameter: 0.9896533678727819
Iteration 7
    Energy: -108.55897868904235
    Norm of gradient: 0.029779155791534718
    Regularization hyperparameter: 0.0022842583349248526
    Variation hyperparameter: 0.9896658420137407
Iteration 8
    Energy: -108.5597807742408
    Norm of gradient: 0.026806610997395394
    Regularization hyperparameter: 0.0015969432334668372
    Variation hyperparameter: 0.9896679708269941
Iteration 9
    Energy: -108.56029897432758
    Norm of gradient: 0.022396314007747324
    Regularization hyperparameter: 0.0009592757735594265
    Variation hyperparameter: 0.9896528106118656