How to simulate the local unitary cluster Jastrow (LUCJ) ansatz

In this guide, we show how to use ffsim to simulate the local unitary cluster Jastrow (LUCJ) ansatz. We’ll use it to calculate an approximation to the ground state energy of an ethene molecule.

First, let’s build the molecule.

import pyscf
import pyscf.mcscf

import ffsim

# Build an ethene molecule
bond_distance = 1.339
a = 0.5 * bond_distance
b = a + 0.5626
c = 0.9289
mol = pyscf.gto.Mole()
mol.build(
    atom=[
        ["C", (0, 0, a)],
        ["C", (0, 0, -a)],
        ["H", (0, c, b)],
        ["H", (0, -c, b)],
        ["H", (0, c, -b)],
        ["H", (0, -c, -b)],
    ],
    basis="sto-6g",
    symmetry="d2h",
)

# Define active space
active_space = range(mol.nelectron // 2 - 2, mol.nelectron // 2 + 2)

# 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
mol_data.run_fci()

print(f"norb = {norb}")
print(f"nelec = {nelec}")

converged SCF energy = -77.8266321248745
Parsing /tmp/tmp5fiwjlii
converged SCF energy = -77.8266321248744
CASCI E = -77.8742165643863  E(CI) = -4.02122442107772  S^2 = 0.0000000
norb = 4
nelec = (2, 2)

Overwritten attributes  get_ovlp get_hcore  of <class 'pyscf.scf.hf_symm.SymAdaptedRHF'>
/home/runner/work/ffsim/ffsim/.tox/docs/lib/python3.13/site-packages/pyscf/gto/mole.py:1293: UserWarning: Function mol.dumps drops attribute energy_nuc because it is not JSON-serializable
  warnings.warn(msg)
/home/runner/work/ffsim/ffsim/.tox/docs/lib/python3.13/site-packages/pyscf/gto/mole.py:1293: UserWarning: Function mol.dumps drops attribute intor_symmetric because it is not JSON-serializable
  warnings.warn(msg)

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) = -77.87421536374032  E_corr = -0.04758323886585372
Energy at initialization: -77.87160024816285

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: 88
  message: CONVERGENCE: RELATIVE REDUCTION OF F <= FACTR*EPSMCH
  success: True
   status: 0
      fun: -77.87387001934972
        x: [-1.152e+00 -8.820e-05 ...  8.723e-07 -3.935e-05]
      nit: 5
      jac: [-8.527e-06  3.837e-05 ...  2.700e-05  2.416e-05]
     nfev: 623
     njev: 7
 hess_inv: <88x88 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: 62
  message: CONVERGENCE: RELATIVE REDUCTION OF F <= FACTR*EPSMCH
  success: True
   status: 0
      fun: -77.87363426200139
        x: [-1.152e+00 -1.636e-04 ...  1.249e-05 -8.832e-05]
      nit: 6
      jac: [-2.984e-05 -3.268e-05 ...  0.000e+00 -4.263e-06]
     nfev: 567
     njev: 9
 hess_inv: <62x62 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: 62
 message: Convergence: Relative reduction of objective function <= ftol.
 success: True
     fun: -77.87421592167836
       x: [-1.267e+00  9.889e-02 ...  2.671e-02 -2.150e-02]
     nit: 9
     jac: [ 5.064e-06 -4.244e-06 ... -1.282e-05 -8.923e-08]
    nfev: 1608
    njev: 10
  nlinop: 988

Iteration 1
    Energy: -77.8736299397149
    Norm of gradient: 0.0012391517375424615
    Regularization hyperparameter: 0.0015872302736352074
    Variation hyperparameter: 0.9987145946018284
Iteration 2
    Energy: -77.87363437571274
    Norm of gradient: 9.976184099796519e-05
    Regularization hyperparameter: 0.0034887597051413348
    Variation hyperparameter: 0.9983182472916641
Iteration 3
    Energy: -77.87370517539703
    Norm of gradient: 0.004386159509693792
    Regularization hyperparameter: 5.417937233503614e-05
    Variation hyperparameter: 0.9981200790749285
Iteration 4
    Energy: -77.87371218001499
    Norm of gradient: 0.0036159554990625904
    Regularization hyperparameter: 1.0147755181369196
    Variation hyperparameter: 0.9981199028377431
Iteration 5
    Energy: -77.87407307284465
    Norm of gradient: 0.010529863250826053
    Regularization hyperparameter: 0.0017374318034521765
    Variation hyperparameter: 0.9960640311483865
Iteration 6
    Energy: -77.87419717822713
    Norm of gradient: 0.0017795927293646915
    Regularization hyperparameter: 0.0005941239281817304
    Variation hyperparameter: 0.9960939810046952
Iteration 7
    Energy: -77.87421037969345
    Norm of gradient: 0.002470601085777093
    Regularization hyperparameter: 0.0003972953478935006
    Variation hyperparameter: 0.9960868398335759
Iteration 8
    Energy: -77.8742153611771
    Norm of gradient: 0.0007007471677068659
    Regularization hyperparameter: 0.00018481429615762492
    Variation hyperparameter: 0.9960817373559102
Iteration 9
    Energy: -77.87421592167836
    Norm of gradient: 0.00012279499963759353
    Regularization hyperparameter: 0.0001846912426612833
    Variation hyperparameter: 0.9960817350003409