Mejora una estimación de SQD con optimización de orbitales

La diagonalización cuántica basada en muestras (SQD) aproxima la energía del estado fundamental diagonalizando el hamiltoniano en un subespacio fijo de configuraciones electrónicas. Esa estimación depende de la base de orbitales en la que se expresa el hamiltoniano, y la optimización de orbitales (OO) aprovecha esta libertad para reducir la energía sin ampliar el subespacio.

Esta guía ejecuta SQD sobre una molécula de $N_2$ y luego mejora el resultado con optimización de orbitales, usando ffsim para representar el hamiltoniano y encontrar la rotación de orbitales que minimiza la energía.

Ejecutar SQD

Construimos las integrales moleculares de $N_2$ en la base de orbitales moleculares (MO), generamos muestras aleatorias uniformes y ejecutamos SQD para obtener una aproximación del estado fundamental.

# Added by doQumentation — installs the packages this notebook needs if they are missing
import importlib.util

_needed = {"ffsim": "ffsim", "numpy": "numpy", "pyscf": "pyscf", "qiskit_addon_sqd": "qiskit-addon-sqd"}
_missing = [pip for module, pip in _needed.items()
            if importlib.util.find_spec(module) is None]
# One at a time, so a package that fails to install does not block the others
for _pip in _missing:
    %pip install -q {_pip}
if not _missing:
    print("\u2713 All required packages are installed")

import numpy as np
import pyscf
import pyscf.cc
import pyscf.mcscf
from qiskit_addon_sqd.counts import generate_bit_array_uniform
from qiskit_addon_sqd.fermion import diagonalize_fermionic_hamiltonian

# Specify molecule properties
num_orbitals = 16
num_elec_a = num_elec_b = 5
spin_sq = 0

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

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

# Get molecular integrals
scf = pyscf.scf.RHF(mol).run()
num_orbitals = len(active_space)
n_electrons = int(sum(scf.mo_occ[active_space]))
num_elec_a = (n_electrons + mol.spin) // 2
num_elec_b = (n_electrons - mol.spin) // 2
cas = pyscf.mcscf.CASCI(scf, num_orbitals, (num_elec_a, num_elec_b))
mo = cas.sort_mo(active_space, base=0)
hcore, nuclear_repulsion_energy = cas.get_h1cas(mo)
eri = pyscf.ao2mo.restore(1, cas.get_h2cas(mo), num_orbitals)

# Compute exact energy
exact_energy = cas.run().e_tot

# Create a seed to control randomness throughout this workflow
rng = np.random.default_rng(24)

# Generate random samples
bit_array = generate_bit_array_uniform(
    10_000, num_orbitals * 2, rand_seed=rng
)

# Run SQD
result = diagonalize_fermionic_hamiltonian(
    hcore,
    eri,
    bit_array,
    samples_per_batch=100,
    norb=num_orbitals,
    nelec=(num_elec_a, num_elec_b),
    num_batches=1,
    max_iterations=5,
    symmetrize_spin=True,
    seed=rng,
)

converged SCF energy = -108.835236570775
CASCI E = -109.046671778080  E(CI) = -32.8155692383187  S^2 = 0.0000000

sqd_energy = result.energy + nuclear_repulsion_energy
print(f"Exact energy:  {exact_energy:.8f}")
print(f"SQD energy:    {sqd_energy:.8f}")

Exact energy:  -109.04667178
SQD energy:    -108.98469255

Optimizar los orbitales

La optimización de orbitales busca una rotación de orbitales que reduzca la energía variacional

$$E = \langle \psi | \mathcal{U}^\dagger\, H\, \mathcal{U} | \psi \rangle$$

de la aproximación del estado fundamental de SQD $|\psi\rangle$. Una rotación de orbitales se especifica mediante una matriz unitaria $N \times N$ $\mathbf{U}$ ($N$ es el número de orbitales espaciales), que actúa sobre el estado de muchos cuerpos a través del operador

$$\mathcal{U} = \exp\left[\sum_{pq, \sigma} \log(\mathbf{U})_{pq}\, a^\dagger_{p\sigma} a_{q\sigma}\right].$$

ffsim.optimize_orbitals devuelve la matriz $\mathbf{U}$, y aplicarla a la base de orbitales (mediante hamiltonian.rotated) equivale a aplicar $\mathcal{U}$ al estado. Consulta la explicación de la rotación de orbitales de ffsim para más detalles.

Dado que rotar los orbitales cambia el hamiltoniano que ve el subespacio, alternamos dos pasos hasta que la energía deja de mejorar:

1. Diagonalizar el hamiltoniano en la base actual sobre el conjunto fijo de configuraciones.

2. Optimizar los orbitales encontrando la rotación que minimiza la energía del estado resultante y, después, rotar las integrales a la nueva base.

Delegamos el paso de rotación de orbitales en ffsim.optimize_orbitals, que encuentra la rotación que minimiza la energía a partir de las matrices de densidad reducidas (RDM) monocuerpo y bicuerpo del estado. Consulta Sec. II A 4 para más detalles.

Por qué la optimización de orbitales ayuda aquí

La base de orbitales moleculares (MO) de SCF es estacionaria respecto a las rotaciones de orbitales para el problema de CI completa. Pero SQD trabaja en un pequeño subespacio truncado (aquí unos pocos cientos de cadenas CI de entre aproximadamente 19 millones de determinantes de CI completa), para el cual la base de MO generalmente no es óptima, de modo que rotar los orbitales reduce la energía que el subespacio puede representar.

import ffsim
from pyscf import fci

# ffsim's ``MolecularHamiltonian`` uses the same "chemist" ordering for the two-body
# tensor as PySCF's ``eri``, and stores the nuclear repulsion energy as the constant
# term so that expectation values come out as total energies.

Alternar diagonalización y optimización de orbitales

Mantenemos el subespacio de diagonalización fijo en las configuraciones descubiertas por SQD anteriormente, de modo que cada iteración aísle el efecto de rotar los orbitales. Cada iteración:

1. Diagonaliza el hamiltoniano sobre el subespacio fijo en la base actual, usando el solucionador de CI seleccionada de PySCF.

2. Construye las RDM del estado resultante, que es todo lo que necesita ffsim.optimize_orbitals.

3. Optimiza los orbitales: ffsim.optimize_orbitals devuelve la rotación que minimiza la energía, que aplicamos a las integrales para pasar a la base mejorada.

Registramos la energía antes de cada paso de optimización. Como la base mejora en cada iteración, esta secuencia disminuye monótonamente hacia la mejor energía alcanzable en el subespacio fijo.

# Fix the diagonalization subspace to the configurations found by SQD.
ci_strings = (result.sci_state.ci_strs_a, result.sci_state.ci_strs_b)
nelec = (num_elec_a, num_elec_b)

# Start from the MO basis in which we ran SQD.
hamiltonian_opt = ffsim.MolecularHamiltonian(
    hcore, eri, constant=nuclear_repulsion_energy
)

num_iters = 10
for i in range(num_iters):
    # Diagonalize over the fixed subspace in the current basis.
    myci = fci.selected_ci.SelectedCI()
    myci = fci.addons.fix_spin_(myci, ss=spin_sq)
    _, amplitudes = fci.selected_ci.kernel_fixed_space(
        myci,
        hamiltonian_opt.one_body_tensor,
        hamiltonian_opt.two_body_tensor,
        num_orbitals,
        nelec,
        ci_strs=ci_strings,
    )

    # Build the RDMs and record the energy before re-optimizing the orbitals.
    dm1, dm2 = myci.make_rdm12(amplitudes, num_orbitals, nelec)
    rdm = ffsim.ReducedDensityMatrix(dm1, dm2)
    energy = rdm.expectation(hamiltonian_opt).real
    print(f"Iteration {i}: energy = {energy:.8f}")

    # Rotate the Hamiltonian into the energy-minimizing basis for the next iteration.
    # optimize_orbitals returns the unitary matrix U minimizing
    # rdm.rotated(U).expectation(hamiltonian), equivalently
    # rdm.expectation(hamiltonian.rotated(U.conj().T)), so we rotate by U^dagger.
    orbital_rotation = ffsim.optimize_orbitals(rdm, hamiltonian_opt)
    hamiltonian_opt = hamiltonian_opt.rotated(orbital_rotation.T.conj())

Iteration 0: energy = -108.98452447
Iteration 1: energy = -108.99981993
Iteration 2: energy = -109.00585329
Iteration 3: energy = -109.00816569
Iteration 4: energy = -109.00936616
Iteration 5: energy = -109.01014322
Iteration 6: energy = -109.01069439
Iteration 7: energy = -109.01109308
Iteration 8: energy = -109.01138928
Iteration 9: energy = -109.01161411

Comparar los resultados

La optimización de orbitales mejora la estimación del subespacio fijo, cerrando buena parte de la brecha con la energía exacta sin llegar a quedar por debajo de ella.

# Diagonalize once more in the final optimized basis to report the improved energy.
myci = fci.selected_ci.SelectedCI()
myci = fci.addons.fix_spin_(myci, ss=spin_sq)
_, amplitudes = fci.selected_ci.kernel_fixed_space(
    myci,
    hamiltonian_opt.one_body_tensor,
    hamiltonian_opt.two_body_tensor,
    num_orbitals,
    nelec,
    ci_strs=ci_strings,
)
dm1, dm2 = myci.make_rdm12(amplitudes, num_orbitals, nelec)
energy_after_oo = (
    ffsim.ReducedDensityMatrix(dm1, dm2).expectation(hamiltonian_opt).real
)

print(f"Exact energy:      {exact_energy:.8f}")
print(f"SQD energy (MO):   {sqd_energy:.8f}")
print(f"Energy after OO:   {energy_after_oo:.8f}")

Exact energy:      -109.04667178
SQD energy (MO):   -108.98469255
Energy after OO:   -109.01178727