How to simulate entanglement forging

In this guide, we show how to simulate entanglement forging.

Build a molecule

We’ll use, as an example, a water molecule at equilibrium bond length in an active space of 6 orbitals and 8 electrons.

import math

import pyscf
import pyscf.mcscf

import ffsim

# Build a water molecule
radius_1 = 0.958
radius_2 = 0.958
bond_angle_deg = 104.478

h1_x = radius_1
h2_x = radius_2 * math.cos(math.pi / 180 * bond_angle_deg)
h2_y = radius_2 * math.sin(math.pi / 180 * bond_angle_deg)

mol = pyscf.gto.Mole()
mol.build(
    atom=[
        ["O", (0, 0, 0)],
        ["H", (h1_x, 0, 0)],
        ["H", (h2_x, h2_y, 0)],
    ],
    basis="sto-6g",
    symmetry="c2v",
)

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

# Get molecular data and molecular 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

# Compute the FCI energy
mol_data.run_fci()

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

converged SCF energy = -75.6787887956297
Parsing /tmp/tmp87spfcj1
converged SCF energy = -75.6787887956314
CASCI E = -75.7288249991515  E(CI) = -23.6332495815006  S^2 = 0.0000000
norb = 6
nelec = (4, 4)

Overwritten attributes  get_hcore get_ovlp  of <class 'pyscf.scf.hf.RHF'>
/home/runner/work/ffsim/ffsim/.tox/docs/lib/python3.12/site-packages/pyscf/gto/mole.py:1294: 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.12/site-packages/pyscf/gto/mole.py:1294: UserWarning: Function mol.dumps drops attribute intor_symmetric because it is not JSON-serializable
  warnings.warn(msg)

Initialize ansatz operator

For our ansatz, we’ll use a dense “brickwork” pattern of “hop gates.” This is implemented in ffsim as the HopGateAnsatzOperator class.

import numpy as np


# Define the pattern of two-qubit gates to use
def brickwork(norb: int, n_layers: int):
    for i in range(n_layers):
        for j in range(i % 2, norb - 1, 2):
            yield (j, j + 1)


n_layers = norb
interaction_pairs = list(brickwork(norb, n_layers))

# Generate random initial parameters
rng = np.random.default_rng(1234)
thetas = rng.normal(scale=1e-1, size=len(interaction_pairs))

# Construct the ansatz operator
operator = ffsim.HopGateAnsatzOperator(norb, interaction_pairs, thetas)

Choose reference occupations

The next step is to choose the reference occupations to use. We’ll use 3 reference occupations, and these will be constructed using the same occupations for the alpha and beta spin orbitals.

# Specify the reference occupations as "spatial" since we'll use the same occupations
# for both alpha and beta spin orbitals
reference_occupations_spatial = [(0, 1, 2, 3), (1, 2, 3, 4), (0, 1, 2, 4)]

# Construct the full reference occupations, including alpha and beta parts
reference_occupations = list(
    zip(reference_occupations_spatial, reference_occupations_spatial)
)

Compute energy

The entanglement forging energy is computed using the multireference_state_prod function. The prod suffix refers to the fact that our ansatz operator is a product operator which acts on the alpha and beta spin sectors independently. In this case, we choose to use the same ansatz operator for both spin sectors, so we pass a tuple (operator, operator) of the same ansatz operator repeated twice.

# Compute the energy of the ansatz
energy, _ = ffsim.multireference_state_prod(
    mol_hamiltonian, (operator, operator), reference_occupations, norb=norb, nelec=nelec
)

print(f"Energy at initialialization: {energy}")

Energy at initialialization: -75.67794403659727

Optimize energy

Now that we know how to initialize an ansatz operator and compute its energy, we can put this logic in a function and minimize the function. Here, we minimizze the function using Scipy’s optimization module.

import scipy.optimize


def fun(x):
    # Initialize the ansatz operator from the parameter vector
    operator = ffsim.HopGateAnsatzOperator(norb, interaction_pairs, x)
    # Compute energy
    energy, _ = ffsim.multireference_state_prod(
        mol_hamiltonian,
        (operator, operator),
        reference_occupations,
        norb=norb,
        nelec=nelec,
    )
    return energy


result = scipy.optimize.minimize(
    fun, x0=operator.thetas, method="L-BFGS-B", options=dict(maxfun=100)
)

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

Number of parameters: 15
  message: STOP: TOTAL NO. of f AND g EVALUATIONS EXCEEDS LIMIT
  success: False
   status: 1
      fun: -75.6838156586247
        x: [-1.603e-01  6.417e-03 ...  5.747e-02 -1.005e-01]
      nit: 3
      jac: [ 2.146e-04  1.080e-04 ... -4.746e-03  7.401e-03]
     nfev: 112
     njev: 7
 hess_inv: <15x15 LbfgsInvHessProduct with dtype=float64>