{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "frontmatter",
      "metadata": {},
      "source": [
        "---\n",
        "title: \"Improve an SQD estimate with orbital optimization\"\n",
        "description: \"Improve an SQD estimate with orbital optimization for the latest version of Sample-based quantum diagonalization (SQD)\"\n",
        "---\n",
        "\n"
      ]
    },
    {
      "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.\n",
        "\n"
      ]
    },
    {
      "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.\n",
        "\n"
      ]
    },
    {
      "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(\n",
        "    10_000, num_orbitals * 2, rand_seed=rng\n",
        ")\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.\n",
        "\n"
      ]
    },
    {
      "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.\n",
        "\n"
      ]
    },
    {
      "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.\n",
        "\n"
      ]
    },
    {
      "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(\n",
        "    hcore, eri, constant=nuclear_repulsion_energy\n",
        ")\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.\n",
        "\n"
      ]
    },
    {
      "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 = (\n",
        "    ffsim.ReducedDensityMatrix(dm1, dm2).expectation(hamiltonian_opt).real\n",
        ")\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}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "id": "a1b8767d",
      "source": "© IBM Corp., 2017-2026"
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}