Create a variational ansatz¶
In this tutorial, we’ll implement the unitary cluster Jastrow (UCJ) ansatz, a variational quantum circuit ansatz for molecular quantum chemistry. ffsim provides a built-in implementation in UCJOpSpinBalanced; in this tutorial we’ll implement our own simplified version by defining a variational operator class. See The local unitary cluster Jastrow (LUCJ) ansatz for a detailed explanation of the ansatz and its properties.
Build the molecule¶
We begin by building a nitrogen molecule in an active space of 6 orbitals and 6 electrons. We also compute the full configuration interaction (FCI) energy to use as a reference for gauging the quality of our variational state.
[1]:
import pyscf
import pyscf.mcscf
import ffsim
# Build N2 molecule
mol = pyscf.gto.Mole()
mol.build(
atom=[["N", (0, 0, -0.5)], ["N", (0.0, 0, 0.5)]],
basis="sto-6g",
symmetry="Dooh",
)
# Define active space
n_frozen = 4
active_space = range(n_frozen, mol.nao_nr())
# Get molecular data and Hamiltonian
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 the FCI (exact) energy for comparison
cas = pyscf.mcscf.CASCI(scf, ncas=norb, nelecas=nelec).run()
fci_energy = cas.e_tot
print(f"norb = {norb}")
print(f"nelec = {nelec}")
print(f"FCI energy = {fci_energy}")
converged SCF energy = -108.464957764795
CASCI E = -108.566842251942 E(CI) = -11.9110176528507 S^2 = 0.0000000
norb = 6
nelec = (3, 3)
FCI energy = -108.5668422519418
The UCJ ansatz¶
The UCJ ansatz prepares a variational quantum state by applying \(L\) layers of operators to a reference state \(|\Phi_0\rangle\), typically the Hartree-Fock state:
The number of layers \(L\) is called the number of ansatz repetitions. Each layer consists of two operations applied in sequence:
An orbital rotation \(\mathcal{U}_k\), specified by a unitary matrix \(U_k\). An orbital rotation transforms the single-particle basis of the state and can be implemented efficiently on a quantum computer.
A Jastrow factor \(e^{i \mathcal{J}_k}\), where \(\mathcal{J}_k\) is a diagonal Coulomb operator
The diagonal Coulomb operator is diagonal in the computational basis, meaning it applies phases to basis states determined solely by the electron occupancies. Conjugating by the orbital rotation \(\mathcal{U}_k\) allows the interaction structure to be varied flexibly in each layer.
For a closed-shell molecule, we use the spin-balanced variant of the UCJ ansatz, in which \(J^{(\alpha\alpha)} = J^{(\beta\beta)}\) and \(J^{(\alpha\beta)} = J^{(\beta\alpha)}\). Each layer is then described by two symmetric matrices \(J^{(\alpha\alpha)}\) and \(J^{(\alpha\beta)}\), along with a single unitary \(U_k\) applied to both spin sectors.
Implementing the operator class¶
ffsim uses a protocol-based design (see Protocols) to allow user-defined classes to integrate with its functions. To make our UCJ operator work with ffsim.apply_unitary, we implement the SupportsApplyUnitary protocol by defining a _apply_unitary_ method with the signature
def _apply_unitary_(
self, vec: np.ndarray, norb: int, nelec: int | tuple[int, int], copy: bool
) -> np.ndarray:
...
Inside _apply_unitary_, we implement each UCJ layer as three steps:
Apply \(\mathcal{U}_k^\dagger\) using apply_orbital_rotation.
Apply the Jastrow factor \(e^{i \mathcal{J}_k}\) using apply_diag_coulomb_evolution. Note that
apply_diag_coulomb_evolutioncomputes \(e^{-it\mathcal{J}}\), so we passtime=-1.0to get \(e^{i\mathcal{J}}\). For the spin-balanced case, \(J^{(\alpha\alpha)} = J^{(\beta\beta)}\), so we pass the same matrix for both.Apply \(\mathcal{U}_k\) using apply_orbital_rotation.
To support variational optimization with a standard optimizer, we also implement from_parameters and to_parameters methods for conversion between the operator and a real-valued parameter vector. For the orbital rotations, we use the parametrization \(U = \exp(K)\) where \(K\) is anti-Hermitian (\(K^\dagger = -K\)). Any anti-Hermitian \(N \times N\) matrix has \(N^2\) real degrees of freedom:
The \(N(N-1)/2\) real off-diagonal entries of the strict upper triangle of \(K\).
The \(N(N+1)/2\) imaginary parts of the upper triangle of \(K\) (including the diagonal, which must be purely imaginary).
For each symmetric diagonal Coulomb matrix, we store the \(N(N+1)/2\) upper-triangle entries. The total parameter count per layer is therefore \(N^2 + 2 \cdot N(N+1)/2 = N^2 + N(N+1)\).
[2]:
from dataclasses import dataclass
import numpy as np
import scipy.linalg
def _unitary_from_params(params: np.ndarray, norb: int) -> np.ndarray:
"""Construct a unitary matrix from a real-valued parameter vector.
Uses the parametrization U = exp(K) where K is anti-Hermitian (K† = -K).
The norb**2 parameters encode K as:
- the first norb*(norb-1)//2 entries: real parts of the strict upper
triangle of K,
- the remaining norb*(norb+1)//2 entries: imaginary parts of the upper
triangle of K.
"""
n_triu_offdiag = norb * (norb - 1) // 2
mat = np.zeros((norb, norb), dtype=complex)
# Imaginary part, on the upper triangle including the diagonal. Assign this
# first, because the real part below touches only the strict upper triangle.
rows, cols = np.triu_indices(norb)
vals = 1j * params[n_triu_offdiag:]
mat[rows, cols] = vals
mat[cols, rows] = vals
# Real part, on the strict upper triangle, antisymmetric
rows, cols = np.triu_indices(norb, k=1)
vals = params[:n_triu_offdiag]
mat[rows, cols] += vals
mat[cols, rows] -= vals
return scipy.linalg.expm(mat)
def _params_from_unitary(unitary: np.ndarray) -> np.ndarray:
"""Extract parameters from a unitary matrix via its matrix logarithm.
For a unitary matrix the logarithm is anti-Hermitian up to numerical error,
so its upper triangle can be read off directly.
"""
norb = unitary.shape[0]
n_triu_offdiag = norb * (norb - 1) // 2
mat = scipy.linalg.logm(unitary)
rows, cols = np.triu_indices(norb, k=1)
params = np.empty(norb**2)
params[:n_triu_offdiag] = mat[rows, cols].real
rows, cols = np.triu_indices(norb)
params[n_triu_offdiag:] = mat[rows, cols].imag
return params
@dataclass
class UCJOp:
"""Spin-balanced UCJ ansatz operator.
Attributes:
diag_coulomb_mats: The diagonal Coulomb matrices, as a NumPy array of
shape ``(n_reps, 2, norb, norb)``. The second axis indexes the spin
interaction type: index 0 for alpha-alpha, index 1 for alpha-beta.
orbital_rotations: The orbital rotations, as a NumPy array of shape
``(n_reps, norb, norb)``.
"""
diag_coulomb_mats: np.ndarray # shape: (n_reps, 2, norb, norb)
orbital_rotations: np.ndarray # shape: (n_reps, norb, norb)
@property
def norb(self) -> int:
"""The number of spatial orbitals."""
return self.diag_coulomb_mats.shape[-1]
@property
def n_reps(self) -> int:
"""The number of ansatz repetitions."""
return self.diag_coulomb_mats.shape[0]
def _apply_unitary_(
self,
vec: np.ndarray,
norb: int,
nelec: int | tuple[int, int],
copy: bool,
) -> np.ndarray:
if isinstance(nelec, int):
return NotImplemented
if copy:
vec = vec.copy()
for (mat_aa, mat_ab), orbital_rotation in zip(
self.diag_coulomb_mats, self.orbital_rotations
):
# Step 1: Apply U†
vec = ffsim.apply_orbital_rotation(
vec, orbital_rotation.conj().T, norb=norb, nelec=nelec, copy=False
)
# Step 2: Apply exp(iJ)
# apply_diag_coulomb_evolution computes exp(-itJ), so time=-1.0
# gives exp(iJ)
# For spin-balanced: J^aa = J^bb, so we pass the same matrix for both
vec = ffsim.apply_diag_coulomb_evolution(
vec,
(mat_aa, mat_ab, mat_aa),
time=-1.0,
norb=norb,
nelec=nelec,
copy=False,
)
# Step 3: Apply U
vec = ffsim.apply_orbital_rotation(
vec, orbital_rotation, norb=norb, nelec=nelec, copy=False
)
return vec
@staticmethod
def n_params(norb: int, n_reps: int) -> int:
"""Return the number of parameters of the ansatz."""
return n_reps * (norb**2 + norb * (norb + 1))
@staticmethod
def from_parameters(params: np.ndarray, *, norb: int, n_reps: int) -> "UCJOp":
"""Initialize a UCJOp from a real-valued parameter vector.
The parameter vector encodes the operator layer by layer. For each of
the ``n_reps`` layers, the layout is:
- norb**2 parameters for the orbital rotation,
- norb*(norb+1)//2 parameters for the alpha-alpha diagonal Coulomb matrix,
- norb*(norb+1)//2 parameters for the alpha-beta diagonal Coulomb matrix.
"""
triu_rows, triu_cols = np.triu_indices(norb)
n_triu = len(triu_rows)
diag_coulomb_mats = np.zeros((n_reps, 2, norb, norb))
orbital_rotations = np.zeros((n_reps, norb, norb), dtype=complex)
index = 0
for k in range(n_reps):
orbital_rotations[k] = _unitary_from_params(
params[index : index + norb**2], norb
)
index += norb**2
for spin_index in range(2):
diag_coulomb_mats[k, spin_index][triu_rows, triu_cols] = params[
index : index + n_triu
]
diag_coulomb_mats[k, spin_index][triu_cols, triu_rows] = params[
index : index + n_triu
]
index += n_triu
return UCJOp(
diag_coulomb_mats=diag_coulomb_mats, orbital_rotations=orbital_rotations
)
def to_parameters(self) -> np.ndarray:
"""Convert the operator to a real-valued parameter vector."""
norb = self.norb
triu_rows, triu_cols = np.triu_indices(norb)
n_triu = len(triu_rows)
params = np.zeros(UCJOp.n_params(norb, self.n_reps))
index = 0
for k in range(self.n_reps):
params[index : index + norb**2] = _params_from_unitary(
self.orbital_rotations[k]
)
index += norb**2
for spin_index in range(2):
params[index : index + n_triu] = self.diag_coulomb_mats[k, spin_index][
triu_rows, triu_cols
]
index += n_triu
return params
Let’s verify the implementation with randomly initialized parameters. We generate a random parameter vector, construct the operator using from_parameters, and apply it to the Hartree-Fock state using ffsim.apply_unitary.
[3]:
rng = np.random.default_rng(1234)
n_reps = 2
# Initialize with a random parameter vector
x = rng.standard_normal(UCJOp.n_params(norb, n_reps))
operator = UCJOp.from_parameters(x, norb=norb, n_reps=n_reps)
# Prepare the Hartree-Fock reference state and Hamiltonian linear operator
reference_state = ffsim.hartree_fock_state(norb, nelec)
hamiltonian_linop = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)
# Apply the operator using the protocol function
ansatz_state = ffsim.apply_unitary(reference_state, operator, norb=norb, nelec=nelec)
energy = np.real(np.vdot(ansatz_state, hamiltonian_linop @ ansatz_state))
print(f"Energy with random parameters: {energy:.6f}")
print(f"FCI energy: {fci_energy:.6f}")
Energy with random parameters: -106.029938
FCI energy: -108.566842
Random parameters generally produce a state with high energy. A better strategy is to initialize the parameters from the results of a classical calculation.
Initializing from CCSD amplitudes¶
A natural way to initialize the UCJ ansatz is from the coupled cluster singles and doubles (CCSD) \(t\)-amplitudes. As explained in The LUCJ ansatz, the connection is established by a double factorization of the \(T_2\) operator:
The Trotter approximation \(e^{T_2 - T_2^\dagger} \approx \prod_k \mathcal{U}_k e^{i\mathcal{J}_k} \mathcal{U}_k^\dagger\) then yields an operator in UCJ form. We use ffsim.linalg.double_factorized_t2 to compute the factorization.
First, let’s run CCSD to obtain the \(t\)-amplitudes.
[4]:
from pyscf import cc
# Run CCSD
ccsd = cc.CCSD(
scf, frozen=[i for i in range(mol.nao_nr()) if i not in active_space]
).run()
# Double-factorize the t2 amplitudes
diag_coulomb_mats_df, orbital_rotations_df = ffsim.linalg.double_factorized_t2(ccsd.t2)
n_vecs = len(orbital_rotations_df)
print(f"CCSD energy: {ccsd.e_tot:.6f}")
print(f"Number of terms from double factorization: {n_vecs}")
E(CCSD) = -108.5658290955831 E_corr = -0.1008713307875628
CCSD energy: -108.565829
Number of terms from double factorization: 18
The double factorization produces n_vecs terms, each with one orbital rotation and one diagonal Coulomb matrix. We’ll truncate to n_reps = 2 layers. In the spin-balanced approximation, the same matrix is used for both the alpha-alpha and alpha-beta interactions, so we stack it to get the shape (n_reps, 2, norb, norb) expected by our class.
Since our class is a plain dataclass, we can also construct it directly from arrays, without going through the parameter vector.
[5]:
n_reps = 2
# Truncate to n_reps layers and stack the same matrix for both spin interactions
orbital_rotations_init = orbital_rotations_df[:n_reps]
diag_coulomb_mats_init = np.stack(
[diag_coulomb_mats_df[:n_reps], diag_coulomb_mats_df[:n_reps]], axis=1
)
# Construct the operator directly from arrays
operator_init = UCJOp(
diag_coulomb_mats=diag_coulomb_mats_init,
orbital_rotations=orbital_rotations_init,
)
# Apply and compute energy
ansatz_state = ffsim.apply_unitary(
reference_state, operator_init, norb=norb, nelec=nelec
)
energy_init = np.real(np.vdot(ansatz_state, hamiltonian_linop @ ansatz_state))
print(f"Energy at CCSD initialization: {energy_init:.6f}")
print(f"FCI energy: {fci_energy:.6f}")
print(f"Error from FCI: {energy_init - fci_energy:.6f}")
Energy at CCSD initialization: -108.552820
FCI energy: -108.566842
Error from FCI: 0.014023
The CCSD-initialized ansatz already gives a much lower energy than random initialization. We can improve it further with variational optimization.
Variational optimization¶
We use to_parameters to extract the initial parameter vector from the CCSD-initialized operator, then define an objective function that calls from_parameters to reconstruct the operator at each optimizer step.
[6]:
import scipy.optimize
# Get the initial parameter vector from the CCSD-initialized operator
x0 = operator_init.to_parameters()
print(f"Number of parameters: {len(x0)}")
def energy_from_params(x: np.ndarray) -> float:
operator = UCJOp.from_parameters(x, norb=norb, n_reps=n_reps)
state = ffsim.apply_unitary(reference_state, operator, norb=norb, nelec=nelec)
return np.real(np.vdot(state, hamiltonian_linop @ state))
result = scipy.optimize.minimize(
energy_from_params,
x0=x0,
method="L-BFGS-B",
options=dict(maxiter=50),
)
print(f"Energy at initialization: {energy_init:.6f}")
print(f"Energy after optimization: {result.fun:.6f}")
print(f"FCI energy: {fci_energy:.6f}")
print(f"Error from FCI: {result.fun - fci_energy:.6f}")
Number of parameters: 156
Energy at initialization: -108.552820
Energy after optimization: -108.561738
FCI energy: -108.566842
Error from FCI: 0.005104
Comparing with the built-in implementation¶
Let’s verify that our UCJOp produces the same result as ffsim’s built-in UCJOpSpinBalanced when given the same arrays.
[7]:
# Apply our custom UCJOp
state_custom = ffsim.apply_unitary(
reference_state, operator_init, norb=norb, nelec=nelec
)
# Apply the built-in UCJOpSpinBalanced with the same arrays
builtin_operator = ffsim.UCJOpSpinBalanced(
diag_coulomb_mats=diag_coulomb_mats_init,
orbital_rotations=orbital_rotations_init,
)
state_builtin = ffsim.apply_unitary(
reference_state, builtin_operator, norb=norb, nelec=nelec
)
overlap = abs(np.vdot(state_custom, state_builtin))
print(f"Overlap between custom and built-in states: {overlap:.15f}")
Overlap between custom and built-in states: 0.999999999999999
The two implementations produce identical states, confirming that our class correctly captures the UCJ ansatz.
The built-in class also provides a from_t_amplitudes method that performs the whole CCSD initialization in a single call: it double-factorizes the \(t_2\) amplitudes, truncates the decomposition to n_reps terms, and stacks each diagonal Coulomb matrix for the two spin interaction types. Let’s check that it reproduces the operator we constructed by hand.
[8]:
# Initialize the built-in operator directly from the CCSD t2 amplitudes
builtin_operator_t2 = ffsim.UCJOpSpinBalanced.from_t_amplitudes(ccsd.t2, n_reps=n_reps)
print(
"Diagonal Coulomb matrices match: "
f"{np.allclose(builtin_operator_t2.diag_coulomb_mats, diag_coulomb_mats_init)}"
)
print(
"Orbital rotations match: "
f"{np.allclose(builtin_operator_t2.orbital_rotations, orbital_rotations_init)}"
)
# Compute the energy of the state prepared by the built-in operator
state_builtin_t2 = ffsim.apply_unitary(
reference_state, builtin_operator_t2, norb=norb, nelec=nelec
)
energy_builtin_t2 = np.real(
np.vdot(state_builtin_t2, hamiltonian_linop @ state_builtin_t2)
)
print(f"Energy from from_t_amplitudes: {energy_builtin_t2:.6f}")
print(f"Energy at CCSD initialization: {energy_init:.6f}")
Diagonal Coulomb matrices match: True
Orbital rotations match: True
Energy from from_t_amplitudes: -108.552820
Energy at CCSD initialization: -108.552820
from_t_amplitudes reproduces our hand-built initialization exactly. On top of the truncation we performed here, it supports several features that we did not implement, including initializing a final orbital rotation from the \(t_1\) amplitudes, restricting the diagonal Coulomb interactions to a chosen set of orbital pairs (giving the local UCJ, or LUCJ, ansatz), and variationally optimizing the double factorization to reduce its error.
You’ve made it to the end of this tutorial! For next steps, consider exploring:
The local unitary cluster Jastrow (LUCJ) ansatz for a deeper treatment of the ansatz, including locality constraints and CCSD initialization.
Simulate the variational quantum eigensolver (VQE) for a guide on running VQE with ffsim, including the linear method optimizer.