Efficiently simulate the single-layer UCJ ansatz

This guide shows how to simulate a single-layer (L)UCJ ansatz using the backpropagation algorithm described in Belagali et al. (2026). The algorithm returns the expectation value

\[\bra{\Psi_0} U^\dagger H U \ket{\Psi_0}\]

where \(\ket{\Psi_0}\) is an electronic configuration (Slater determinant), \(U\) is the UCJ ansatz operator, and \(H\) is a molecular Hamiltonian. Note that the algorithm does not support sampling from the state \(U \ket{\Psi_0}\), and it cannot simulate UCJ operators with more than one layer. The running time of the algorithm is \(O(N^7)\), where \(N\) is the number of orbitals.

import numpy as np
import pyscf
import pyscf.cc
import scipy

import ffsim

LUCJ ansatz for a closed-shell molecule

We construct an LUCJ ansatz for a nitrogen molecule in the 6-31G basis set. Since it’s a closed-shell system, we use the spin-balanced UCJ ansatz.

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

# Define active space
n_frozen = 2
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, nelec = mol_data.norb, mol_data.nelec
mol_hamiltonian = mol_data.hamiltonian
print(f"norb = {norb}")
print(f"nelec = {nelec}")

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

n_reps = 1

# Define interactions
pairs_aa = [(p, p + 1) for p in range(norb - 1)]
pairs_ab = None

ucj_op = ffsim.UCJOpSpinBalanced.from_t_amplitudes(
    ccsd.t2,
    t1=ccsd.t1,
    n_reps=n_reps,
    interaction_pairs=(pairs_aa, pairs_ab),
    optimize=True,
    regularization=0.1,
    options=dict(maxiter=100),
)

converged SCF energy = -108.835236570774
norb = 16
nelec = (5, 5)
E(CCSD) = -109.0398256929733  E_corr = -0.2045891221988312

Now, let’s compute the energy of the ansatz state using the backpropagation method, which is implemented in the ffsim.ucj_energy_spin_balanced function. We’ll also compute the energy using direct state vector simulation and check that the results match.

energy = ffsim.ucj_energy_spin_balanced(ucj_op, mol_hamiltonian, nelec)

# Compute energy from state vector simulation and check that it matches
reference_state = ffsim.hartree_fock_state(norb, nelec)
vec = ffsim.apply_unitary(reference_state, ucj_op, norb=norb, nelec=nelec)
linop = ffsim.linear_operator(mol_hamiltonian, norb=norb, nelec=nelec)
energy_state_vec = np.vdot(vec, linop @ vec)
np.testing.assert_allclose(energy, energy_state_vec)

print(f"HF energy: {scf.e_tot:.6f}")
print(f"LUCJ Energy: {energy:.6f}")
print(f"CCSD energy: {ccsd.e_tot:.6f}")

HF energy: -108.835237
LUCJ Energy: -108.906690
CCSD energy: -109.039826

In addition to functions that directly return the energy, ffsim also implements function factories for returning both the energy and the gradient, using JAX for automatic differentiation. This is useful for variationally optimizing the ansatz parameters to lower the energy. The following code block demonstrates this using the ucj_energy_and_grad_func_spin_balanced function, whose output can be passed directly to scipy.optimize.minimize. The optional chunk_size argument sets the maximum number of two-body Hamiltonian terms to evaluate at a time. Smaller values use less memory but take longer to run, so setting it can help when running on GPUs with limited memory.

def callback(intermediate_result: scipy.optimize.OptimizeResult):
    print(f" {intermediate_result.fun:.6f}")


value_and_grad = ffsim.ucj_energy_and_grad_func_spin_balanced(
    mol_hamiltonian,
    nelec,
    interaction_pairs=(pairs_aa, pairs_ab),
    with_final_orbital_rotation=True,
    chunk_size=65536,
)
result = scipy.optimize.minimize(
    value_and_grad,
    ucj_op.to_parameters(interaction_pairs=(pairs_aa, pairs_ab)),
    method="L-BFGS-B",
    jac=True,
    callback=callback,
    options=dict(maxiter=5),
)
optimal_ucj = ffsim.UCJOpSpinBalanced.from_parameters(
    result.x,
    norb=norb,
    n_reps=n_reps,
    interaction_pairs=(pairs_aa, pairs_ab),
    with_final_orbital_rotation=True,
)
energy = ffsim.ucj_energy_spin_balanced(optimal_ucj, mol_hamiltonian, nelec)
print(f"Optimized LUCJ energy: {energy:.6f}")

 -108.952858
 -108.963896
 -108.969690
 -108.970705
 -108.971477
Optimized LUCJ energy: -108.971477

UCJ ansatz for an open-shell molecule

Now, we demonstrate the method applied to an open-shell molecule, where we use the spin-unbalanced UCJ ansatz. The molecule here is a hydroxyl radical in the 6-31G basis set. For this example, we’ll use unrestricted pair connectivity, i.e. the full UCJ ansatz.

# Build HO molecule
mol = pyscf.gto.Mole()
mol.build(
    atom=[["H", (0, 0, 0)], ["O", (0, 0, 1.0)]],
    basis="6-31g",
    spin=1,
    # ROHF fails to converge with Coov, so use the C2v subgroup
    symmetry="C2v",
)

# Get molecular data and Hamiltonian
scf = pyscf.scf.ROHF(mol).run()
mol_data = ffsim.MolecularData.from_scf(scf)
norb, nelec = mol_data.norb, mol_data.nelec
mol_hamiltonian = mol_data.hamiltonian
print(f"norb = {norb}")
print(f"nelec = {nelec}")

# Get CCSD t2 amplitudes for initializing the ansatz
ccsd = pyscf.cc.CCSD(scf).run()

ucj_op = ffsim.UCJOpSpinUnbalanced.from_t_amplitudes(
    ccsd.t2,
    t1=ccsd.t1,
    n_reps=1,
    optimize=True,
    regularization=0.1,
    options=dict(maxiter=100),
)

energy = ffsim.ucj_energy_spin_unbalanced(ucj_op, mol_hamiltonian, nelec)
print(f"HF energy: {scf.e_tot:.6f}")
print(f"UCJ Energy: {energy:.6f}")
print(f"CCSD energy: {ccsd.e_tot:.6f}")

value_and_grad = ffsim.ucj_energy_and_grad_func_spin_unbalanced(
    mol_hamiltonian,
    nelec,
    with_final_orbital_rotation=ucj_op.final_orbital_rotation is not None,
)
result = scipy.optimize.minimize(
    value_and_grad,
    ucj_op.to_parameters(),
    method="L-BFGS-B",
    jac=True,
    callback=callback,
    options=dict(maxiter=5),
)
optimal_ucj = ffsim.UCJOpSpinUnbalanced.from_parameters(
    result.x,
    norb=norb,
    n_reps=1,
    with_final_orbital_rotation=ucj_op.final_orbital_rotation is not None,
)
energy = ffsim.ucj_energy_spin_unbalanced(optimal_ucj, mol_hamiltonian, nelec)
print(f"Optimized UCJ energy: {energy:.6f}")

converged SCF energy = -75.3608533074516
norb = 11
nelec = (5, 4)

WARN: RCCSD method does not support ROHF method. ROHF object is converted to UHF object and UCCSD method is called.

E(UCCSD) = -75.46256067414772  E_corr = -0.1017073666961219
HF energy: -75.360853
UCJ Energy: -75.401747
CCSD energy: -75.462561
 -75.435085
 -75.437992
 -75.440204
 -75.441313
 -75.442150
Optimized UCJ energy: -75.442150