Add unit tests for configuration recovery module (#48)

* Add tests for post_select_by_hamming_weight

* Add unit tests for configuration_recovery module

qiskit_addon_sqd/configuration_recovery.py

@@ -43,15 +43,15 @@ def post_select_by_hamming_weight(
         hamming_left: The target hamming weight of the left half of bitstrings
 
     Returns:
-        A mask signifying which samples were selected from the input matrix.
+        A mask signifying which samples (rows) were selected from the input matrix.
     """
     if hamming_left < 0 or hamming_right < 0:
         raise ValueError("Hamming weights must be non-negative integers.")
     num_bits = bitstring_matrix.shape[1]
 
     # Find the bitstrings with correct hamming weight on both halves
-    up_keepers = np.sum(bitstring_matrix[:, : num_bits // 2], axis=1) == hamming_left
-    down_keepers = np.sum(bitstring_matrix[:, num_bits // 2 :], axis=1) == hamming_right
+    up_keepers = np.sum(bitstring_matrix[:, num_bits // 2 :], axis=1) == hamming_right
+    down_keepers = np.sum(bitstring_matrix[:, : num_bits // 2], axis=1) == hamming_left
     correct_bs_mask = np.array(np.logical_and(up_keepers, down_keepers))
 
     return correct_bs_mask
@@ -68,9 +68,20 @@ def recover_configurations(
     """
     Refine bitstrings based on average orbital occupancy and a target hamming weight.
 
-    This function makes the assumption that bit ``i`` represents the spin-down orbital
-    corresponding to the spin-up orbital in bit ``i + N`` where ``N`` is the number of
-    spatial orbitals and ``i < N``.
+    This function refines each bit in isolation in an attempt to transform the Hilbert space
+    represented by the input ``bitstring_matrix`` into a space closer to that which supports
+    the ground state.
+
+    .. note::
+
+        This function makes the assumption that bit ``i`` represents the spin-down orbital
+        corresponding to the spin-up orbital in bit ``i + N`` where ``N`` is the number of
+        spatial orbitals and ``i < N``.
+
+    .. note::
+
+        The output configurations may not necessarily have correct hamming weight, as each bit
+        is flipped in isolation from the other bits in the bitstring.
 
     Args:
         bitstring_matrix: A 2D array of ``bool`` representations of bit
@@ -84,9 +95,11 @@ def recover_configurations(
         rand_seed: A seed to control random behavior
 
     Returns:
-        A refined bitstring matrix and an updated probability array. The bitstrings in
-        the output matrix may still have incorrect hamming weight, but in the aggregate, it is
-        hoped the samples are higher quality.
+        A refined bitstring matrix and an updated probability array.
+
+    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].
     """
     if num_elec_a < 0 or num_elec_b < 0:
         raise ValueError("The numbers of electrons must be specified as non-negative integers.")
@@ -111,7 +124,7 @@ def recover_configurations(
     return bs_mat_out, freqs_out
 
 
-def _p_flip_0_to_1(ratio_exp: float, occ: float, eps: float = 0.01) -> float:
+def _p_flip_0_to_1(ratio_exp: float, occ: float, eps: float = 0.01) -> float:  # pragma: no cover
     """
     Calculate the probability of flipping a bit from 0 to 1.
 
@@ -135,12 +148,14 @@ def _p_flip_0_to_1(ratio_exp: float, occ: float, eps: float = 0.01) -> float:
     # Occupancy is >= naive expectation.
     # The probability to flip the bit increases linearly from ``eps`` to
     # ``~1.0`` as the occupation deviates further from the expected ratio
+    if ratio_exp == 1.0:
+        return eps
     slope = (1 - eps) / (1 - ratio_exp)
     intercept = 1 - slope
     return occ * slope + intercept
 
 
-def _p_flip_1_to_0(ratio_exp: float, occ: float, eps: float = 0.01) -> float:
+def _p_flip_1_to_0(ratio_exp: float, occ: float, eps: float = 0.01) -> float:  # pragma: no cover
     """
     Calculate the probability of flipping a bit from 1 to 0.
 
@@ -165,6 +180,8 @@ def _p_flip_1_to_0(ratio_exp: float, occ: float, eps: float = 0.01) -> float:
 
     # Occupancy is >= naive expectation.
     # Flip 1s to 0 with small (~eps) probability in this case
+    if ratio_exp == 0.0:
+        return 1 - eps
     slope = -eps / ratio_exp
     intercept = eps / ratio_exp
     return occ * slope + intercept

test/test_configuration_recovery.py

@@ -14,7 +14,94 @@
 
 import unittest
 
+import numpy as np
+import pytest
+from qiskit_addon_sqd.configuration_recovery import (
+    post_select_by_hamming_weight,
+    recover_configurations,
+)
+
 
 class TestConfigurationRecovery(unittest.TestCase):
-    def test_configuration_recovery(self):
-        pass
+    def setUp(self):
+        self.small_mat = np.array(
+            [[False, False, True, False, False, False], [False, False, True, True, False, False]]
+        )
+
+    def test_post_select_by_hamming_weight(self):
+        with self.subTest("Empty test"):
+            ham_l = 1
+            ham_r = 0
+            empty_mat = np.empty((0, 6))
+            bs_mask = post_select_by_hamming_weight(
+                empty_mat, hamming_right=ham_r, hamming_left=ham_l
+            )
+            self.assertEqual(0, bs_mask.size)
+        with self.subTest("Basic test"):
+            ham_l = 1
+            ham_r = 0
+            expected = np.array([True, False])
+            bs_mask = post_select_by_hamming_weight(
+                self.small_mat, hamming_right=ham_r, hamming_left=ham_l
+            )
+            self.assertTrue((expected == bs_mask).all())
+        with self.subTest("Bad hamming"):
+            ham_l = 0
+            ham_r = -1
+            with pytest.raises(ValueError) as e_info:
+                post_select_by_hamming_weight(
+                    self.small_mat, hamming_right=ham_r, hamming_left=ham_l
+                )
+            assert e_info.value.args[0] == "Hamming weights must be non-negative integers."
+
+    def test_recover_configurations(self):
+        with self.subTest("Empty test"):
+            num_orbs = 6
+            ham_l = 1
+            ham_r = 0
+            empty_mat = np.empty((0, num_orbs))
+            empty_probs = np.empty((0,))
+            occs = [False] * num_orbs
+            mat_rec, probs_rec = recover_configurations(
+                empty_mat, empty_probs, occs, num_elec_a=ham_r, num_elec_b=ham_l
+            )
+            self.assertEqual(0, mat_rec.size)
+            self.assertEqual(0, probs_rec.size)
+        with self.subTest("Basic test. Zeros to ones."):
+            bs_mat = np.array([[False, False, False, False]])
+            probs = np.array([1.0])
+            occs = [1.0, 1.0, 1.0, 1.0]
+            num_a = 2
+            num_b = 2
+            expected_mat = np.array([[True, True, True, True]])
+            expected_probs = np.array([1.0])
+            mat_rec, probs_rec = recover_configurations(
+                bs_mat, probs, occs, num_a, num_b, rand_seed=4224
+            )
+            self.assertTrue((expected_mat == mat_rec).all())
+            self.assertTrue((expected_probs == probs_rec).all())
+        with self.subTest("Basic test. Ones to zeros."):
+            bs_mat = np.array([[True, True, True, True]])
+            probs = np.array([1.0])
+            occs = [0.0, 0.0, 0.0, 0.0]
+            num_a = 0
+            num_b = 0
+            expected_mat = np.array([[False, False, False, False]])
+            expected_probs = np.array([1.0])
+            mat_rec, probs_rec = recover_configurations(
+                bs_mat, probs, occs, num_a, num_b, rand_seed=4224
+            )
+            self.assertTrue((expected_mat == mat_rec).all())
+            self.assertTrue((expected_probs == probs_rec).all())
+        with self.subTest("Bad hamming."):
+            bs_mat = np.array([[True, True, True, True]])
+            probs = np.array([1.0])
+            occs = [0.0, 0.0, 0.0, 0.0]
+            num_a = 0
+            num_b = -1
+            with pytest.raises(ValueError) as e_info:
+                recover_configurations(bs_mat, probs, occs, num_a, num_b, rand_seed=4224)
+            assert (
+                e_info.value.args[0]
+                == "The numbers of electrons must be specified as non-negative integers."
+            )