Refactor how mean orbital occupancies are specified throughout API (#151)

* Flip orbital occupancies inside the solver.

* Make type right

* style

* Remove flip function from docs files

* oo

* update reno

* Return 1 1D array for the occupancies

* peer review

* Fix notebooks

* typo

* mypy

* lint

Author: caleb-johnson

docs/apidocs/qiskit_addon_sqd.fermion.rst

@@ -12,7 +12,6 @@ Fermion (:mod:`qiskit_addon_sqd.fermion`)
 .. autoclass:: SCIState
 .. autofunction:: bitstring_matrix_to_ci_strs
 .. autofunction:: enlarge_batch_from_transitions
-.. autofunction:: flip_orbital_occupancies
 .. autofunction:: solve_fermion
 .. autofunction:: optimize_orbitals
 .. autofunction:: rotate_integrals

docs/conf.py

@@ -122,7 +122,6 @@
     ("qiskit_addon_sqd.counts", "normalize_counts_dict"),
     ("qiskit_addon_sqd.fermion", "bitstring_matrix_to_ci_strs"),
     ("qiskit_addon_sqd.fermion", "enlarge_batch_from_transitions"),
-    ("qiskit_addon_sqd.fermion", "flip_orbital_occupancies"),
     ("qiskit_addon_sqd.fermion", "solve_fermion"),
     ("qiskit_addon_sqd.fermion", "optimize_orbitals"),
     ("qiskit_addon_sqd.fermion", "rotate_integrals"),

docs/how_tos/choose_subspace_dimension.ipynb

@@ -139,7 +139,6 @@
     "from qiskit_addon_sqd.configuration_recovery import recover_configurations\n",
     "from qiskit_addon_sqd.fermion import (\n",
     "    bitstring_matrix_to_ci_strs,\n",
-    "    flip_orbital_occupancies,\n",
     "    solve_fermion,\n",
     ")\n",
     "from qiskit_addon_sqd.subsampling import postselect_and_subsample\n",
@@ -162,13 +161,13 @@
     "    e_hist = np.zeros((iterations, n_batches))  # energy history\n",
     "    s_hist = np.zeros((iterations, n_batches))  # spin history\n",
     "    d_hist = np.zeros((iterations, n_batches))  # subspace dimension history\n",
-    "    occupancy_hist = np.zeros((iterations, 2 * num_orbitals))\n",
-    "    occupancies_bitwise = None  # orbital i corresponds to column i in bitstring matrix\n",
+    "    occupancy_hist = []\n",
+    "    avg_occupancy = None\n",
     "    for i in range(iterations):\n",
     "        print(f\"Starting configuration recovery iteration {i}\")\n",
     "        # On the first iteration, we have no orbital occupancy information from the\n",
     "        # solver, so we just post-select from the full bitstring set based on hamming weight.\n",
-    "        if occupancies_bitwise is None:\n",
+    "        if avg_occupancy is None:\n",
     "            bs_mat_tmp = bitstring_matrix_full\n",
     "            probs_arr_tmp = probs_arr_full\n",
     "\n",
@@ -178,7 +177,7 @@
     "            bs_mat_tmp, probs_arr_tmp = recover_configurations(\n",
     "                bitstring_matrix_full,\n",
     "                probs_arr_full,\n",
-    "                occupancies_bitwise,\n",
+    "                avg_occupancy,\n",
     "                num_elec_a,\n",
     "                num_elec_b,\n",
     "                rand_seed=rng,\n",
@@ -199,7 +198,7 @@
     "        int_e = np.zeros(n_batches)\n",
     "        int_s = np.zeros(n_batches)\n",
     "        int_d = np.zeros(n_batches)\n",
-    "        int_occs = np.zeros((n_batches, 2 * num_orbitals))\n",
+    "        int_occs = []\n",
     "        cs = []\n",
     "        for j in range(n_batches):\n",
     "            ci_strs = bitstring_matrix_to_ci_strs(batches[j], open_shell=open_shell)\n",
@@ -215,20 +214,17 @@
     "            energy_sci += nuclear_repulsion_energy\n",
     "            int_e[j] = energy_sci\n",
     "            int_s[j] = spin\n",
-    "            int_occs[j, :num_orbitals] = avg_occs[0]\n",
-    "            int_occs[j, num_orbitals:] = avg_occs[1]\n",
+    "            int_occs.append(avg_occs)\n",
     "            cs.append(coeffs_sci)\n",
     "\n",
     "        # Combine batch results\n",
-    "        avg_occupancy = np.mean(int_occs, axis=0)\n",
-    "        # The occupancies from the solver should be flipped to match the bits in the bitstring matrix.\n",
-    "        occupancies_bitwise = flip_orbital_occupancies(avg_occupancy)\n",
+    "        avg_occupancy = tuple(np.mean(int_occs, axis=0))\n",
     "\n",
     "        # Track optimization history\n",
     "        e_hist[i, :] = int_e\n",
     "        s_hist[i, :] = int_s\n",
     "        d_hist[i, :] = int_d\n",
-    "        occupancy_hist[i, :] = avg_occupancy\n",
+    "        occupancy_hist.append(avg_occupancy)\n",
     "\n",
     "    return e_hist.flatten(), d_hist.flatten()"
    ]

docs/how_tos/integrate_dice_solver.ipynb

@@ -29,9 +29,6 @@
     "from qiskit_addon_dice_solver import solve_fermion\n",
     "from qiskit_addon_sqd.configuration_recovery import recover_configurations\n",
     "from qiskit_addon_sqd.counts import counts_to_arrays, generate_counts_uniform\n",
-    "from qiskit_addon_sqd.fermion import (\n",
-    "    flip_orbital_occupancies,\n",
-    ")\n",
     "from qiskit_addon_sqd.subsampling import postselect_and_subsample\n",
     "\n",
     "# Specify molecule properties\n",
@@ -87,12 +84,12 @@
     "samples_per_batch = 300\n",
     "\n",
     "# Self-consistent configuration recovery loop\n",
-    "occupancies_bitwise = None  # orbital i corresponds to column i in bitstring matrix\n",
     "e_hist = np.zeros((iterations, n_batches))\n",
+    "avg_occupancy = None\n",
     "for i in range(iterations):\n",
     "    # On the first iteration, we have no orbital occupancy information from the\n",
     "    # solver, so we just post-select from the full bitstring set based on hamming weight.\n",
-    "    if occupancies_bitwise is None:\n",
+    "    if avg_occupancy is None:\n",
     "        bs_mat_tmp = bitstring_matrix_full\n",
     "        probs_arr_tmp = probs_arr_full\n",
     "\n",
@@ -102,7 +99,7 @@
     "        bs_mat_tmp, probs_arr_tmp = recover_configurations(\n",
     "            bitstring_matrix_full,\n",
     "            probs_arr_full,\n",
-    "            occupancies_bitwise,\n",
+    "            avg_occupancy,\n",
     "            num_elec_a,\n",
     "            num_elec_b,\n",
     "            rand_seed=rng,\n",
@@ -120,7 +117,7 @@
     "    )\n",
     "    # Run eigenstate solvers in a loop. This loop should be parallelized for larger problems.\n",
     "    int_e = np.zeros(n_batches)\n",
-    "    int_occs = np.zeros((n_batches, 2 * num_orbitals))\n",
+    "    int_occs = []\n",
     "    for j, batch in enumerate(batches):\n",
     "        energy_sci, wf_mags, avg_occs = solve_fermion(\n",
     "            batch,\n",
@@ -129,13 +126,10 @@
     "            mpirun_options=[\"-n\", \"8\"],\n",
     "        )\n",
     "        int_e[j] = energy_sci + nuclear_repulsion_energy\n",
-    "        int_occs[j, :num_orbitals] = avg_occs[0]\n",
-    "        int_occs[j, num_orbitals:] = avg_occs[1]\n",
+    "        int_occs.append(avg_occs)\n",
     "\n",
     "    # Combine batch results\n",
-    "    avg_occupancy = np.mean(int_occs, axis=0)\n",
-    "    # The occupancies from the solver should be flipped to match the bits in the bitstring matrix.\n",
-    "    occupancies_bitwise = flip_orbital_occupancies(avg_occupancy)\n",
+    "    avg_occupancy = tuple(np.mean(int_occs, axis=0))\n",
     "\n",
     "    # Track optimization history\n",
     "    e_hist[i, :] = int_e"

docs/how_tos/use_oo_to_optimize_hamiltonian_basis.ipynb

@@ -141,7 +141,7 @@
    ],
    "source": [
     "from qiskit_addon_sqd.configuration_recovery import recover_configurations\n",
-    "from qiskit_addon_sqd.fermion import flip_orbital_occupancies, solve_fermion\n",
+    "from qiskit_addon_sqd.fermion import solve_fermion\n",
     "from qiskit_addon_sqd.subsampling import postselect_and_subsample\n",
     "\n",
     "# SQSD options\n",
@@ -155,13 +155,13 @@
     "# Self-consistent configuration recovery loop\n",
     "e_hist = np.zeros((iterations, n_batches))  # energy history\n",
     "s_hist = np.zeros((iterations, n_batches))  # spin history\n",
-    "occupancy_hist = np.zeros((iterations, 2 * num_orbitals))\n",
-    "occupancies_bitwise = None  # orbital i corresponds to column i in bitstring matrix\n",
+    "occupancy_hist = []\n",
+    "avg_occupancy = None\n",
     "for i in range(iterations):\n",
     "    print(f\"Starting configuration recovery iteration {i}\")\n",
     "    # On the first iteration, we have no orbital occupancy information from the\n",
     "    # solver, so we just post-select from the full bitstring set based on hamming weight.\n",
-    "    if occupancies_bitwise is None:\n",
+    "    if avg_occupancy is None:\n",
     "        bs_mat_tmp = bitstring_matrix_full\n",
     "        probs_arr_tmp = probs_arr_full\n",
     "\n",
@@ -171,7 +171,7 @@
     "        bs_mat_tmp, probs_arr_tmp = recover_configurations(\n",
     "            bitstring_matrix_full,\n",
     "            probs_arr_full,\n",
-    "            occupancies_bitwise,\n",
+    "            avg_occupancy,\n",
     "            num_elec_a,\n",
     "            num_elec_b,\n",
     "            rand_seed=rng,\n",
@@ -191,7 +191,7 @@
     "    # Run eigenstate solvers in a loop. This loop should be parallelized for larger problems.\n",
     "    int_e = np.zeros(n_batches)\n",
     "    int_s = np.zeros(n_batches)\n",
-    "    int_occs = np.zeros((n_batches, 2 * num_orbitals))\n",
+    "    int_occs = []\n",
     "    cs = []\n",
     "    for j in range(n_batches):\n",
     "        energy_sci, coeffs_sci, avg_occs, spin = solve_fermion(\n",
@@ -205,19 +205,16 @@
     "        energy_sci += nuclear_repulsion_energy\n",
     "        int_e[j] = energy_sci\n",
     "        int_s[j] = spin\n",
-    "        int_occs[j, :num_orbitals] = avg_occs[0]\n",
-    "        int_occs[j, num_orbitals:] = avg_occs[1]\n",
+    "        int_occs.append(avg_occs)\n",
     "        cs.append(coeffs_sci)\n",
     "\n",
     "    # Combine batch results\n",
-    "    avg_occupancy = np.mean(int_occs, axis=0)\n",
-    "    # The occupancies from the solver should be flipped to match the bits in the bitstring matrix.\n",
-    "    occupancies_bitwise = flip_orbital_occupancies(avg_occupancy)\n",
+    "    avg_occupancy = tuple(np.mean(int_occs, axis=0))\n",
     "\n",
     "    # Track optimization history\n",
     "    e_hist[i, :] = int_e\n",
     "    s_hist[i, :] = int_s\n",
-    "    occupancy_hist[i, :] = avg_occupancy"
+    "    occupancy_hist.append(avg_occupancy)"
    ]
   },
   {
@@ -350,9 +347,9 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Exact energy: -109.04667177808028\n",
+      "Exact energy: -109.04667177808032\n",
       "SQD energy: -108.7531706601421\n",
-      "Energy after OO: -108.80400806164369\n"
+      "Energy after OO: -108.80400806164377\n"
      ]
     }
    ],

docs/tutorials/01_chemistry_hamiltonian.ipynb

@@ -15,6 +15,7 @@
 
 from __future__ import annotations
 
+import warnings
 from collections import defaultdict
 from collections.abc import Sequence
 
@@ -51,7 +52,7 @@ def post_select_by_hamming_weight(
 def recover_configurations(
     bitstring_matrix: np.ndarray,
     probabilities: Sequence[float],
-    avg_occupancies: np.ndarray,
+    avg_occupancies: tuple[np.ndarray, np.ndarray],
     num_elec_a: int,
     num_elec_b: int,
     rand_seed: np.random.Generator | int | None = None,
@@ -72,9 +73,9 @@ def recover_configurations(
         bitstring_matrix: A 2D array of ``bool`` representations of bit
             values such that each row represents a single bitstring
         probabilities: A 1D array specifying a probability distribution over the bitstrings
-        avg_occupancies: A 1D array containing the mean occupancy of each orbital. It is assumed
-            that ``avg_occupancies[i]`` corresponds to the orbital represented by column
-            ``i`` in ``bitstring_matrix``.
+        avg_occupancies: A length-2 tuple of arrays holding the mean occupancy of the spin-up
+            and spin-down orbitals, respectively. The occupancies should be formatted as:
+            ``(array([occ_a_0, ..., occ_a_N]), array([occ_b_0, ..., occ_b_N]))``
         num_elec_a: The number of spin-up electrons in the system.
         num_elec_b: The number of spin-down electrons in the system.
         rand_seed: A seed for controlling randomness
@@ -85,20 +86,28 @@ def recover_configurations(
     References:
         [1]: J. Robledo-Moreno, et al., `Chemistry Beyond Exact Solutions on a Quantum-Centric Supercomputer <https://arxiv.org/abs/2405.05068>`_,
              arXiv:2405.05068 [quant-ph].
-
     """
     rng = np.random.default_rng(rand_seed)
 
+    occ_dims = len(np.array(avg_occupancies).shape)
+    if occ_dims == 1:
+        warnings.warn(
+            "Passing avg_occupancies as a 1D array is deprecated. Pass a length-2 tuple containing the spin-up and spin-down occupancies respectively.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        norb = bitstring_matrix.shape[1] // 2
+        avg_occupancies = (np.flip(avg_occupancies[norb:]), np.flip(avg_occupancies[:norb]))
+
     if num_elec_a < 0 or num_elec_b < 0:
         raise ValueError("The numbers of electrons must be specified as non-negative integers.")
 
-    # First, we need to flip the orbitals such that
-
     corrected_dict: defaultdict[str, float] = defaultdict(float)
+    occs_array = np.flip(avg_occupancies).flatten()
     for bitstring, freq in zip(bitstring_matrix, probabilities):
         bs_corrected = _bipartite_bitstring_correcting(
             bitstring,
-            avg_occupancies,
+            occs_array,
             num_elec_a,
             num_elec_b,
             rng=rng,

qiskit_addon_sqd/configuration_recovery.py

@@ -79,7 +79,7 @@ def solve_fermion(
     spin_sq: float | None = None,
     max_davidson: int = 100,
     verbose: int | None = None,
-) -> tuple[float, SCIState, list[np.ndarray], float]:
+) -> tuple[float, SCIState, tuple[np.ndarray, np.ndarray], float]:
     """Approximate the ground state given molecular integrals and a set of electronic configurations.
 
     Args:
@@ -107,7 +107,7 @@ def solve_fermion(
     Returns:
         - Minimum energy from SCI calculation
         - The SCI ground state
-        - Average occupancy of the alpha and beta orbitals, respectively
+        - Tuple containing orbital occupancies for spin-up and spin-down orbitals. Formatted as: ``(array([occ_a_0, ..., occ_a_N]), array([occ_b_0, ..., occ_b_N]))``
         - Expectation value of spin-squared
 
     """
@@ -145,7 +145,7 @@ def solve_fermion(
 
     # Calculate the avg occupancy of each orbital
     dm1 = myci.make_rdm1s(sci_vec, norb, (num_up, num_dn))
-    avg_occupancy = [np.diagonal(dm1[0]), np.diagonal(dm1[1])]
+    avg_occupancy = (np.diagonal(dm1[0]), np.diagonal(dm1[1]))
 
     # Compute total spin
     spin_squared = myci.spin_square(sci_vec, norb, (num_up, num_dn))[0]
@@ -175,7 +175,7 @@ def optimize_orbitals(
     num_steps_grad: int = 10_000,
     learning_rate: float = 0.01,
     max_davidson: int = 100,
-) -> tuple[float, np.ndarray, list[np.ndarray]]:
+) -> tuple[float, np.ndarray, tuple[np.ndarray, np.ndarray]]:
     """Optimize orbitals to produce a minimal ground state.
 
     The process involves iterating over 3 steps:
@@ -219,7 +219,7 @@ def optimize_orbitals(
     Returns:
         - The groundstate energy found during the last optimization iteration
         - An optimized 1D array defining the orbital transform
-        - Average orbital occupancy
+        - Tuple containing orbital occupancies for spin-up and spin-down orbitals. Formatted as: ``(array([occ_a_0, ..., occ_a_N]), array([occ_b_0, ..., occ_b_N]))``
 
     """
     norb = hcore.shape[0]
@@ -270,14 +270,15 @@ def optimize_orbitals(
         dm1, dm2_chem = myci.make_rdm12(amplitudes, norb, (num_up, num_dn))
         dm2 = np.asarray(dm2_chem.transpose(0, 2, 3, 1), order="C")
         dm1a, dm1b = myci.make_rdm1s(amplitudes, norb, (num_up, num_dn))
+        avg_occupancy = (np.diagonal(dm1a), np.diagonal(dm1b))
 
         # TODO: Expose the momentum parameter as an input option
         # Optimize the basis rotations
         _optimize_orbitals_sci(
             k_flat, learning_rate, 0.9, num_steps_grad, dm1, dm2, hcore, eri_phys
         )
 
-    return e_qsci, k_flat, [np.diagonal(dm1a), np.diagonal(dm1b)]
+    return e_qsci, k_flat, avg_occupancy
 
 
 def rotate_integrals(
@@ -320,29 +321,6 @@ def rotate_integrals(
     return np.array(hcore_rot), np.array(eri_rot)
 
 
-def flip_orbital_occupancies(occupancies: np.ndarray) -> np.ndarray:
-    """Flip an orbital occupancy array to match the indexing of a bitstring.
-
-    This function reformats a 1D array of spin-orbital occupancies formatted like:
-
-        ``[occ_a_1, occ_a_2, ..., occ_a_N, occ_b_1, ..., occ_b_N]``
-
-    To an array formatted like:
-
-        ``[occ_a_N, ..., occ_a_1, occ_b_N, ..., occ_b_1]``
-
-    where ``N`` is the number of spatial orbitals.
-    """
-    norb = occupancies.shape[0] // 2
-    occ_up = occupancies[:norb]
-    occ_dn = occupancies[norb:]
-    occ_out = np.zeros(2 * norb)
-    occ_out[:norb] = np.flip(occ_up)
-    occ_out[norb:] = np.flip(occ_dn)
-
-    return occ_out
-
-
 @deprecate_func(
     removal_timeline="no sooner than qiskit-addon-sqd 0.10.0",
     since="0.6.0",

qiskit_addon_sqd/fermion.py

@@ -0,0 +1,11 @@
+---
+upgrade:
+  - |
+    Removed :func:`qiskit_addon_sqd.fermion.flip_orbital_occupancies`. Users no longer need to flip the orbital occupancies output from :func:`qiskit_addon_sqd.fermion.solve_fermion` and :func:`qiskit_addon_sqd.fermion.optimize_orbitals`; they will be output in the order expected by :func:`qiskit_addon_sqd.configuration_recovery.recover_configurations`: ``tuple(array([occ_a_0, ..., occ_a_N]), array([occ_b_0, ..., occ_b_N]))``.
+deprecations:
+  - |
+    The ``avg_occupancies`` argument to :func:`qiskit_addon_sqd.configuration_recovery.recover_configurations` should now be a length-2 tuple containing the spin-up and spin-down occupancies, respectively.
+
+    Old format: ``array([occ_b_N, ..., occ_b_0, occ_a_N, ..., occ_a_0])``
+
+    New format: ``tuple(array([occ_a_0, ..., occ_a_N]), array([occ_b_0, ..., occ_b_N]))``