Add the ability to checkpoint configuration recovery

docs/guides/integrate_fulqrum.ipynb

@@ -95,7 +95,7 @@
     "        strs_b = [format(x, f\"0{norb}b\") for x in ints_b]\n",
     "        subspace = fq.Subspace([strs_a, strs_b])\n",
     "        Hsub.update_subspace(subspace)\n",
-    "        Hsub_csr_linop = Hsub.to_csr_linearoperator_fast(verbose=False)\n",
+    "        Hsub_csr_linop = Hsub.to_csr_linearoperator_fast()\n",
     "        diag = Hsub.diagonal_vector()\n",
     "        v0 = np.zeros(len(subspace), dtype=Hsub.dtype)\n",
     "        v0[np.argmin(diag)] = 1.0\n",
@@ -184,7 +184,7 @@
     "        strs_b = [format(x, f\"0{norb}b\") for x in ints_b]\n",
     "        subspace = fq.Subspace([strs_a, strs_b])\n",
     "        Hsub.update_subspace(subspace)\n",
-    "        Hsub_csr_linop = Hsub.to_csr_linearoperator_fast(verbose=False)\n",
+    "        Hsub_csr_linop = Hsub.to_csr_linearoperator_fast()\n",
     "        diag = Hsub.diagonal_vector()\n",
     "        v0 = np.zeros(len(subspace), dtype=Hsub.dtype)\n",
     "        v0[np.argmin(diag)] = 1.0\n",

qiskit_addon_sqd/fermion.py

@@ -225,6 +225,7 @@ def diagonalize_fermionic_hamiltonian(
     carryover_threshold: float = 1e-4,
     callback: Callable[[list[SCIResult]], None] | None = None,
     seed: int | np.random.Generator | None = None,
+    initial_state: SCIState | None = None,
 ) -> SCIResult:
     """Run the sample-based quantum diagonalization (SQD) algorithm.
 
@@ -302,6 +303,12 @@ def diagonalize_fermionic_hamiltonian(
             function, which is a list of (energy, sci_state, occupancies) triplets,
             where each triplet contains the result of a diagonalization.
         seed: A seed for the pseudorandom number generator.
+        initial_state: An ``SCIState`` object from a previous execution used to warm-start
+            the configuration recovery loop. When provided, the algorithm will extract
+            high-weight configurations from this state to use as carryover strings for
+            the first iteration. If ``initial_occupancies`` is not explicitly provided,
+            the initial guess for the average orbital occupancies will also be derived
+            from this state.
 
     Returns:
         The estimate of the energy and the SCI state with that energy.
@@ -380,8 +387,19 @@ def diagonalize_fermionic_hamiltonian(
 
     include_a = np.unique(include_a)
     include_b = np.unique(include_b)
-    carryover_strings_a = np.array([], dtype=np.int64)
-    carryover_strings_b = np.array([], dtype=np.int64)
+
+    if initial_state is not None:
+        # Recover carry over bitstrings from initial_state
+        carryover_strings_a, carryover_strings_b = _extract_carryover_strings(
+            initial_state, carryover_threshold, symmetrize_spin
+        )
+        # Use initial_state orbital occupancies, if occupancies were not
+        # input by the user
+        if current_occupancies is None:
+            current_occupancies = initial_state.orbital_occupancies()
+    else:
+        carryover_strings_a = np.array([], dtype=np.int64)
+        carryover_strings_b = np.array([], dtype=np.int64)
 
     # Convert BitArray into bitstring and probability arrays
     raw_bitstrings, raw_probs = bit_array_to_arrays(bit_array)
@@ -560,6 +578,64 @@ def _prepare_ci_strings(
     return ci_strings
 
 
+def _extract_carryover_strings(
+    sci_state: SCIState,
+    carryover_threshold: float,
+    symmetrize_spin: bool,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Extract high-weight CI strings from an SCIState for configuration carryover.
+
+    Args:
+        sci_state: The SCI state from which to extract the configurations.
+        carryover_threshold: Threshold for carrying over bitstrings with large CI
+            weight from one iteration of configuration recovery to the next.
+            All single-spin CI strings associated with configurations whose coefficient
+            has absolute value greater than this threshold will be extracted.
+        symmetrize_spin: Whether to always merge spin-alpha and spin-beta CI strings
+            into a single list, so that the extracted subspace is invariant with
+            respect to the exchange of spin alpha with spin beta.
+
+    Returns:
+        A tuple ``(carryover_strings_a, carryover_strings_b)`` containing 1D arrays
+        of the extracted spin-alpha and spin-beta CI strings, sorted in descending
+        order by marginal weight.
+    """
+    # Sorted amplitude indices based on absolute values
+    flattened = sci_state.amplitudes.reshape(-1)
+    absolute_vals = np.abs(flattened)
+
+    # Get indices of CI amplitudes greater than or equal to carryover_threshold
+    # np.nonzero is highly optimized in C to find indices matching a condition
+    carryover_indices = np.nonzero(absolute_vals >= carryover_threshold)[0]
+
+    # Extract unique alpha and beta CI bitstrings that participate in configurations
+    # with amplitudes greater than or equal to carryover_threshold
+    _, n_strings_b = sci_state.amplitudes.shape
+    alpha_indices, beta_indices = np.divmod(carryover_indices, n_strings_b)
+    alpha_indices = np.unique(alpha_indices)
+    beta_indices = np.unique(beta_indices)
+    carryover_strings_a = sci_state.ci_strs_a[alpha_indices]
+    carryover_strings_b = sci_state.ci_strs_b[beta_indices]
+
+    # Sort carryover bitstrings in descending order by marginal weight
+    amplitudes_a = sci_state.amplitudes[alpha_indices]
+    weights_a = np.sum(np.square(amplitudes_a.real) + np.square(amplitudes_a.imag), axis=1)
+    amplitudes_b = sci_state.amplitudes[:, beta_indices]
+    weights_b = np.sum(np.square(amplitudes_b.real) + np.square(amplitudes_b.imag), axis=0)
+
+    if symmetrize_spin:
+        carryover_strings = np.concatenate((carryover_strings_a, carryover_strings_b))
+        weights = np.concatenate((weights_a, weights_b))
+        carryover_strings = carryover_strings[np.argsort(weights)[::-1]]
+        carryover_strings = _unique_with_order_preserved(carryover_strings)
+        carryover_strings_a = carryover_strings_b = carryover_strings
+    else:
+        carryover_strings_a = carryover_strings_a[np.argsort(weights_a)[::-1]]
+        carryover_strings_b = carryover_strings_b[np.argsort(weights_b)[::-1]]
+
+    return carryover_strings_a, carryover_strings_b
+
+
 def _process_sci_results(
     config: _LoopConfig,
     results: list[SCIResult],
@@ -606,29 +682,9 @@ def _process_sci_results(
 
     # Carry over bitstrings with large CI weight
     sci_state = current_result.sci_state
-    flattened = sci_state.amplitudes.reshape(-1)
-    absolute_vals = np.abs(flattened)
-    indices = np.argsort(absolute_vals)
-    carryover_index = np.searchsorted(absolute_vals, config.carryover_threshold, sorter=indices)
-    carryover_indices = indices[carryover_index:]
-    _, n_strings_b = sci_state.amplitudes.shape
-    alpha_indices, beta_indices = np.divmod(carryover_indices, n_strings_b)
-    alpha_indices = np.unique(alpha_indices)
-    beta_indices = np.unique(beta_indices)
-    carryover_strings_a = sci_state.ci_strs_a[alpha_indices]
-    carryover_strings_b = sci_state.ci_strs_b[beta_indices]
-    # Sort carryover strings in descending order by marginal weight
-    weights_a = np.sum(np.abs(sci_state.amplitudes[alpha_indices]) ** 2, axis=1)
-    weights_b = np.sum(np.abs(sci_state.amplitudes[:, beta_indices]) ** 2, axis=0)
-    if config.symmetrize_spin:
-        carryover_strings = np.concatenate((carryover_strings_a, carryover_strings_b))
-        weights = np.concatenate((weights_a, weights_b))
-        carryover_strings = carryover_strings[np.argsort(weights)[::-1]]
-        carryover_strings = _unique_with_order_preserved(carryover_strings)
-        carryover_strings_a = carryover_strings_b = carryover_strings
-    else:
-        carryover_strings_a = carryover_strings_a[np.argsort(weights_a)[::-1]]
-        carryover_strings_b = carryover_strings_b[np.argsort(weights_b)[::-1]]
+    carryover_strings_a, carryover_strings_b = _extract_carryover_strings(
+        sci_state, config.carryover_threshold, config.symmetrize_spin
+    )
 
     return _IterationState(
         best_result=best_result,

releasenotes/notes/support-checkpoint-configuration-recovery-0447edc4e24dafa4.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    Added the ability to checkpoint configuration recovery in
+    ``diagonalize_fermionic_hamiltonian`` by providing an ``initial_state`` 
+    (an ``SCIState`` object).

test/test_fermion.py

@@ -341,6 +341,133 @@ def test_diagonalize_fermionic_hamiltonian_reproducible_with_seed(self):
         np.testing.assert_array_equal(result1.sci_state.ci_strs_a, result2.sci_state.ci_strs_a)
         np.testing.assert_allclose(result1.sci_state.ci_strs_b, result2.sci_state.ci_strs_b)
 
+    def test_diagonalize_fermionic_hamiltonian_resume(self):
+        """Test diagonalize_fermionic_hamiltonian restart capability via initial_state."""
+        # Build N2 molecule
+        mol = pyscf.gto.Mole()
+        mol.build(
+            atom=[["N", (0, 0, 0)], ["N", (1.0, 0, 0)]],
+            basis="sto-6g",
+            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()
+        norb = len(active_space)
+        n_electrons = int(sum(scf.mo_occ[active_space]))
+        n_alpha = (n_electrons + mol.spin) // 2
+        n_beta = (n_electrons - mol.spin) // 2
+        nelec = (n_alpha, n_beta)
+        cas = pyscf.mcscf.CASCI(scf, norb, nelec)
+        mo = cas.sort_mo(active_space, base=0)
+        hcore, _ = cas.get_h1cas(mo)
+        eri = pyscf.ao2mo.restore(1, cas.get_h2cas(mo), norb)
+
+        # Generate random bitstrings
+        bit_array = generate_bit_array_uniform(2_000, 2 * norb, rand_seed=self.rng)
+
+        # Define shared keyword arguments
+        kwargs = dict(
+            one_body_tensor=hcore,
+            two_body_tensor=eri,
+            bit_array=bit_array,
+            samples_per_batch=10,
+            norb=norb,
+            nelec=nelec,
+            symmetrize_spin=True,
+            energy_tol=1e-12,
+            occupancies_tol=1e-12,
+            # Note: initial_occupancies is left as default None so the code
+            # extracts it from the initial_state argument on the resumed run
+        )
+
+        # Create two independent but identical Random Number Generators
+        rng_continuous = np.random.default_rng(12345)
+        rng_interrupted = np.random.default_rng(12345)
+
+        # Execute the continuous reference (2 iterations) using the first RNG
+        result_continuous = diagonalize_fermionic_hamiltonian(
+            max_iterations=2, seed=rng_continuous, **kwargs
+        )
+
+        # Execute the interrupted run (1 iteration) using the second RNG
+        result_interrupted = diagonalize_fermionic_hamiltonian(
+            max_iterations=1, seed=rng_interrupted, **kwargs
+        )
+
+        # Resume the interrupted run for 1 more iteration using the same RNG.
+        # Its internal state was stored during the interrupted run, so it
+        # can be resumed and its final state should match the continuous run.
+        result_resumed = diagonalize_fermionic_hamiltonian(
+            max_iterations=1,
+            initial_state=result_interrupted.sci_state,
+            seed=rng_interrupted,
+            **kwargs,
+        )
+
+        # Check that the resumed run's results exactly match the continuous
+        # run's results
+        np.testing.assert_allclose(result_resumed.energy, result_continuous.energy)
+        np.testing.assert_allclose(
+            result_resumed.sci_state.amplitudes, result_continuous.sci_state.amplitudes
+        )
+        np.testing.assert_array_equal(
+            result_resumed.sci_state.ci_strs_a, result_continuous.sci_state.ci_strs_a
+        )
+        np.testing.assert_array_equal(
+            result_resumed.sci_state.ci_strs_b, result_continuous.sci_state.ci_strs_b
+        )
+
+    def test_extract_carryover_strings(self):
+        """Test the extraction and sorting logic of carryover strings."""
+        # Import the private helper directly for testing
+        from qiskit_addon_sqd.fermion import _extract_carryover_strings
+
+        # Construct a highly predictable 3x3 mock state
+        ci_strs_a = np.array([10, 20, 30])
+        ci_strs_b = np.array([1, 2, 3])
+
+        # Amplitudes are chosen so we can easily calculate marginal weights:
+        # Alpha weights (sum of squares across rows): [0.65, 0.36, 0.04]
+        # Beta weights (sum of squares across cols):  [0.37, 0.64, 0.00]
+        amplitudes = np.array([[0.1, 0.8, 0.0], [0.6, 0.0, 0.0], [0.0, 0.0, 0.2]])
+
+        sci_state = SCIState(
+            amplitudes=amplitudes, ci_strs_a=ci_strs_a, ci_strs_b=ci_strs_b, norb=3, nelec=(1, 1)
+        )
+
+        # Test without spin symmetrization
+        # A threshold of 0.5 should only retain amplitudes 0.8 (0, 1)
+        # and 0.6 (1, 0).
+        str_a, str_b = _extract_carryover_strings(
+            sci_state, carryover_threshold=0.5, symmetrize_spin=False
+        )
+
+        # Alpha keeps row indices 0 and 1 (strings 10 and 20).
+        # Their weights are 0.65 and 0.36, so they should be sorted as [10, 20].
+        np.testing.assert_array_equal(str_a, [10, 20])
+
+        # Beta keeps column indices 1 and 0 (strings 2 and 1).
+        # Their weights are 0.64 and 0.37, so they should be sorted as [2, 1].
+        np.testing.assert_array_equal(str_b, [2, 1])
+
+        # Test with spin symmetrization
+        str_a_sym, str_b_sym = _extract_carryover_strings(
+            sci_state, carryover_threshold=0.5, symmetrize_spin=True
+        )
+
+        # Symmetrization merges the kept strings into: [10, 20, 1, 2]
+        # Their respective weights are: [0.65, 0.36, 0.37, 0.64]
+        # Sorted by weight descending: 0.65 (10), 0.64 (2), 0.37 (1), 0.36 (20)
+        expected_sym = [10, 2, 1, 20]
+
+        np.testing.assert_array_equal(str_a_sym, expected_sym)
+        np.testing.assert_array_equal(str_b_sym, expected_sym)
+
     def test_bitstring_matrix_to_ci_strs(self):
         norb = 57
         bitstring = "001111101111111110110001011101100001010000100101100001010"