{ "cells": [ { "cell_type": "markdown", "id": "bb5a576d", "metadata": {}, "source": [ "# Improve an SQD estimate with orbital optimization\n", "\n", "Sample-based quantum diagonalization (SQD) approximates a ground-state energy by\n", "diagonalizing the Hamiltonian in a fixed subspace of electronic configurations. That\n", "estimate depends on the orbital basis in which the Hamiltonian is expressed, and\n", "*orbital optimization* (OO) exploits this freedom to lower the energy without enlarging\n", "the subspace.\n", "\n", "This guide runs SQD on an $N_2$ molecule and then improves the result with orbital\n", "optimization, using [`ffsim`](https://qiskit-community.github.io/ffsim/) to represent\n", "the Hamiltonian and find the energy-minimizing orbital rotation." ] }, { "cell_type": "markdown", "id": "77ab953b", "metadata": {}, "source": [ "## Run SQD\n", "\n", "We build the molecular integrals for $N_2$ in the molecular-orbital (MO) basis, generate\n", "uniform random samples, and run SQD to obtain a ground-state approximation." ] }, { "cell_type": "code", "execution_count": 1, "id": "b8d5618e", "metadata": { "execution": { "iopub.execute_input": "2026-07-16T01:36:48.334261Z", "iopub.status.busy": "2026-07-16T01:36:48.334063Z", "iopub.status.idle": "2026-07-16T01:37:42.585853Z", "shell.execute_reply": "2026-07-16T01:37:42.584303Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "converged SCF energy = -108.835236570775\n", "CASCI E = -109.046671778080 E(CI) = -32.8155692383187 S^2 = 0.0000000\n" ] } ], "source": [ "import numpy as np\n", "import pyscf\n", "import pyscf.cc\n", "import pyscf.mcscf\n", "from qiskit_addon_sqd.counts import generate_bit_array_uniform\n", "from qiskit_addon_sqd.fermion import diagonalize_fermionic_hamiltonian\n", "\n", "# Specify molecule properties\n", "num_orbitals = 16\n", "num_elec_a = num_elec_b = 5\n", "spin_sq = 0\n", "\n", "# Build N2 molecule\n", "mol = pyscf.gto.Mole()\n", "mol.build(\n", " atom=[[\"N\", (0, 0, 0)], [\"N\", (1.0, 0, 0)]],\n", " basis=\"6-31g\",\n", " symmetry=\"Dooh\",\n", ")\n", "\n", "# Define active space\n", "n_frozen = 2\n", "active_space = range(n_frozen, mol.nao_nr())\n", "\n", "# Get molecular integrals\n", "scf = pyscf.scf.RHF(mol).run()\n", "num_orbitals = len(active_space)\n", "n_electrons = int(sum(scf.mo_occ[active_space]))\n", "num_elec_a = (n_electrons + mol.spin) // 2\n", "num_elec_b = (n_electrons - mol.spin) // 2\n", "cas = pyscf.mcscf.CASCI(scf, num_orbitals, (num_elec_a, num_elec_b))\n", "mo = cas.sort_mo(active_space, base=0)\n", "hcore, nuclear_repulsion_energy = cas.get_h1cas(mo)\n", "eri = pyscf.ao2mo.restore(1, cas.get_h2cas(mo), num_orbitals)\n", "\n", "# Compute exact energy\n", "exact_energy = cas.run().e_tot\n", "\n", "# Create a seed to control randomness throughout this workflow\n", "rng = np.random.default_rng(24)\n", "\n", "\n", "# Generate random samples\n", "bit_array = generate_bit_array_uniform(10_000, num_orbitals * 2, rand_seed=rng)\n", "\n", "# Run SQD\n", "result = diagonalize_fermionic_hamiltonian(\n", " hcore,\n", " eri,\n", " bit_array,\n", " samples_per_batch=100,\n", " norb=num_orbitals,\n", " nelec=(num_elec_a, num_elec_b),\n", " num_batches=1,\n", " max_iterations=5,\n", " symmetrize_spin=True,\n", " seed=rng,\n", ")" ] }, { "cell_type": "code", "execution_count": 2, "id": "0ca70028", "metadata": { "execution": { "iopub.execute_input": "2026-07-16T01:37:42.590790Z", "iopub.status.busy": "2026-07-16T01:37:42.589477Z", "iopub.status.idle": "2026-07-16T01:37:42.595563Z", "shell.execute_reply": "2026-07-16T01:37:42.595133Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Exact energy: -109.04667178\n", "SQD energy: -108.98469255\n" ] } ], "source": [ "sqd_energy = result.energy + nuclear_repulsion_energy\n", "print(f\"Exact energy: {exact_energy:.8f}\")\n", "print(f\"SQD energy: {sqd_energy:.8f}\")" ] }, { "cell_type": "markdown", "id": "9160d09c", "metadata": {}, "source": [ "## Optimize the orbitals\n", "\n", "Orbital optimization searches for an orbital rotation that lowers the variational energy\n", "\n", "$$\n", "E = \\langle \\psi | \\mathcal{U}^\\dagger\\, H\\, \\mathcal{U} | \\psi \\rangle\n", "$$\n", "\n", "of the SQD ground-state approximation $|\\psi\\rangle$. An orbital rotation is specified by\n", "an $N \\times N$ unitary matrix $\\mathbf{U}$ ($N$ is the number of spatial orbitals), which\n", "acts on the many-body state through the operator\n", "\n", "$$\n", "\\mathcal{U} = \\exp\\left[\\sum_{pq, \\sigma} \\log(\\mathbf{U})_{pq}\\, a^\\dagger_{p\\sigma} a_{q\\sigma}\\right].\n", "$$\n", "\n", "`ffsim.optimize_orbitals` returns the **matrix** $\\mathbf{U}$, and applying it to the\n", "orbital basis (via `hamiltonian.rotated`) is equivalent to applying $\\mathcal{U}$ to the\n", "state. See the\n", "[ffsim orbital-rotation explanation](https://qiskit-community.github.io/ffsim/explanations/orbital-rotation.html)\n", "for details.\n", "\n", "Since rotating the orbitals changes the Hamiltonian seen by the subspace, we alternate\n", "two steps until the energy stops improving:\n", "\n", "1. **Diagonalize** the Hamiltonian in the current basis over the fixed set of\n", " configurations.\n", "2. **Optimize the orbitals** by finding the rotation that minimizes the energy of the\n", " resulting state, then rotate the integrals into the new basis.\n", "\n", "We delegate the orbital-rotation step to\n", "[`ffsim.optimize_orbitals`](https://qiskit-community.github.io/ffsim/api/ffsim.html#ffsim.optimize_orbitals),\n", "which finds the energy-minimizing rotation from the one- and two-body reduced density\n", "matrices (RDMs) of the state. See\n", "[Sec. II A 4](https://arxiv.org/pdf/2405.05068) for details." ] }, { "cell_type": "markdown", "id": "006745b5", "metadata": {}, "source": [ "### Why orbital optimization helps here\n", "\n", "The SCF molecular-orbital (MO) basis is stationary with respect to orbital rotations\n", "for the *full*-CI problem. But SQD works in a small truncated subspace (here a few\n", "hundred CI strings out of roughly 19 million full-CI determinants), for which the MO\n", "basis is generally not optimal, so rotating the orbitals lowers the energy the subspace\n", "can represent." ] }, { "cell_type": "code", "execution_count": 3, "id": "46eb2f5a", "metadata": { "execution": { "iopub.execute_input": "2026-07-16T01:37:42.597573Z", "iopub.status.busy": "2026-07-16T01:37:42.597399Z", "iopub.status.idle": "2026-07-16T01:37:42.654809Z", "shell.execute_reply": "2026-07-16T01:37:42.654267Z" } }, "outputs": [], "source": [ "import ffsim\n", "from pyscf import fci\n", "\n", "# ffsim's ``MolecularHamiltonian`` uses the same \"chemist\" ordering for the two-body\n", "# tensor as PySCF's ``eri``, and stores the nuclear repulsion energy as the constant\n", "# term so that expectation values come out as total energies." ] }, { "cell_type": "markdown", "id": "93179dc4", "metadata": {}, "source": [ "### Alternate diagonalization and orbital optimization\n", "\n", "We keep the diagonalization subspace **fixed** to the configurations discovered by SQD\n", "above, so that each iteration isolates the effect of rotating the orbitals. Each\n", "iteration:\n", "\n", "1. **Diagonalizes** the Hamiltonian over the fixed subspace in the current basis, using\n", " PySCF's selected-CI solver.\n", "2. **Builds the RDMs** of the resulting state, which is all `ffsim.optimize_orbitals`\n", " needs.\n", "3. **Optimizes the orbitals**: `ffsim.optimize_orbitals` returns the energy-minimizing\n", " rotation, which we apply to the integrals to move into the improved basis.\n", "\n", "We record the energy *before* each optimization step. Because the basis improves every\n", "iteration, this sequence decreases monotonically toward the best energy achievable in\n", "the fixed subspace." ] }, { "cell_type": "code", "execution_count": 4, "id": "a0783e88", "metadata": { "execution": { "iopub.execute_input": "2026-07-16T01:37:42.657394Z", "iopub.status.busy": "2026-07-16T01:37:42.657213Z", "iopub.status.idle": "2026-07-16T01:38:48.374706Z", "shell.execute_reply": "2026-07-16T01:38:48.373697Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Iteration 0: energy = -108.98452447\n", "Iteration 1: energy = -108.99981993\n", "Iteration 2: energy = -109.00585329\n", "Iteration 3: energy = -109.00816569\n", "Iteration 4: energy = -109.00936616\n", "Iteration 5: energy = -109.01014322\n", "Iteration 6: energy = -109.01069439\n", "Iteration 7: energy = -109.01109308\n", "Iteration 8: energy = -109.01138928\n", "Iteration 9: energy = -109.01161411\n" ] } ], "source": [ "# Fix the diagonalization subspace to the configurations found by SQD.\n", "ci_strings = (result.sci_state.ci_strs_a, result.sci_state.ci_strs_b)\n", "nelec = (num_elec_a, num_elec_b)\n", "\n", "# Start from the MO basis in which we ran SQD.\n", "hamiltonian_opt = ffsim.MolecularHamiltonian(hcore, eri, constant=nuclear_repulsion_energy)\n", "\n", "num_iters = 10\n", "for i in range(num_iters):\n", " # Diagonalize over the fixed subspace in the current basis.\n", " myci = fci.selected_ci.SelectedCI()\n", " myci = fci.addons.fix_spin_(myci, ss=spin_sq)\n", " _, amplitudes = fci.selected_ci.kernel_fixed_space(\n", " myci,\n", " hamiltonian_opt.one_body_tensor,\n", " hamiltonian_opt.two_body_tensor,\n", " num_orbitals,\n", " nelec,\n", " ci_strs=ci_strings,\n", " )\n", "\n", " # Build the RDMs and record the energy before re-optimizing the orbitals.\n", " dm1, dm2 = myci.make_rdm12(amplitudes, num_orbitals, nelec)\n", " rdm = ffsim.ReducedDensityMatrix(dm1, dm2)\n", " energy = rdm.expectation(hamiltonian_opt).real\n", " print(f\"Iteration {i}: energy = {energy:.8f}\")\n", "\n", " # Rotate the Hamiltonian into the energy-minimizing basis for the next iteration.\n", " # optimize_orbitals returns the unitary matrix U minimizing\n", " # rdm.rotated(U).expectation(hamiltonian), equivalently\n", " # rdm.expectation(hamiltonian.rotated(U.conj().T)), so we rotate by U^dagger.\n", " orbital_rotation = ffsim.optimize_orbitals(rdm, hamiltonian_opt)\n", " hamiltonian_opt = hamiltonian_opt.rotated(orbital_rotation.T.conj())" ] }, { "cell_type": "markdown", "id": "a61dbda8", "metadata": {}, "source": [ "### Compare the results\n", "\n", "Orbital optimization improves the fixed-subspace estimate, closing much of the gap to\n", "the exact energy while staying above it." ] }, { "cell_type": "code", "execution_count": 5, "id": "762b0903", "metadata": { "execution": { "iopub.execute_input": "2026-07-16T01:38:48.377620Z", "iopub.status.busy": "2026-07-16T01:38:48.377404Z", "iopub.status.idle": "2026-07-16T01:38:53.111168Z", "shell.execute_reply": "2026-07-16T01:38:53.110312Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Exact energy: -109.04667178\n", "SQD energy (MO): -108.98469255\n", "Energy after OO: -109.01178727\n" ] } ], "source": [ "# Diagonalize once more in the final optimized basis to report the improved energy.\n", "myci = fci.selected_ci.SelectedCI()\n", "myci = fci.addons.fix_spin_(myci, ss=spin_sq)\n", "_, amplitudes = fci.selected_ci.kernel_fixed_space(\n", " myci,\n", " hamiltonian_opt.one_body_tensor,\n", " hamiltonian_opt.two_body_tensor,\n", " num_orbitals,\n", " nelec,\n", " ci_strs=ci_strings,\n", ")\n", "dm1, dm2 = myci.make_rdm12(amplitudes, num_orbitals, nelec)\n", "energy_after_oo = ffsim.ReducedDensityMatrix(dm1, dm2).expectation(hamiltonian_opt).real\n", "\n", "print(f\"Exact energy: {exact_energy:.8f}\")\n", "print(f\"SQD energy (MO): {sqd_energy:.8f}\")\n", "print(f\"Energy after OO: {energy_after_oo:.8f}\")" ] } ], "metadata": { "kernelspec": { "display_name": ".venv", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.13.13" } }, "nbformat": 4, "nbformat_minor": 5 }