Fix test.

quantumlib · Jul 21, 2023 · 6fa627e · 6fa627e
1 parent 363257c
commit 6fa627e
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 37 deletions.
diff --git a/src/stim/stabilizers/pauli_string_iter.pybind.cc b/src/stim/stabilizers/pauli_string_iter.pybind.cc
@@ -92,6 +92,25 @@ void stim_pybind::pybind_pauli_string_iter_methods(
         )DOC")
             .data());
 
+    c.def(
+        "iter_next",
+        [](PauliStringIterator<stim::MAX_BITWORD_WIDTH> &self) {
+            return self.iter_next();
+        },
+        clean_doc_string(R"DOC(
+            Get the next permutation of qubit labels.
+
+            It's alot easier to test more complicated edge cases in python
+            which largely arise due to the algorithm for generating the next
+            permutation. The user should not interact with this function.
+
+            Examples:
+                >>> import stim
+                >>> iter = stim.PauliString.iter_all(4, 2, 2)
+                >>> iter.next_qubit_iteration
+        )DOC")
+            .data());
+
     c.def(
         "seed_iterator",
         [](PauliStringIterator<stim::MAX_BITWORD_WIDTH> &self, const pybind11::object &permutation) {

diff --git a/src/stim/stabilizers/pauli_string_iter.test.cc b/src/stim/stabilizers/pauli_string_iter.test.cc
@@ -80,19 +80,17 @@ TEST_EACH_WORD_SIZE_W(pauli_string_iter, iter_pauli_string, {
         iter2.iter_next();
     }
     ASSERT_EQ(iter2.result, PauliString<W>::from_str("+_YYYYY"));
-    // Segfault
     iter2.iter_next();
     // Should be on weight 6 paulis now
-    // ASSERT_EQ(iter2.result, PauliString<W>::from_str("+XXXXXX"));
-    // num_qubits = 4;
-    // min_weight = 3;
-    // max_weight = 4;
-    // PauliStringIterator<W> iter3(num_qubits, min_weight, max_weight);
-    // // std::cout << " XXXX " << std::endl;
-    // for (int i = 0; i < 4 * pow(3, 3) + pow(3, 4); i++) {
-    //     iter3.iter_next();
-    //     // std::cout << iter3.result << std::endl;
-    // }
+    ASSERT_EQ(iter2.result, PauliString<W>::from_str("+XXXXXX"));
+    num_qubits = 4;
+    min_weight = 3;
+    max_weight = 4;
+    PauliStringIterator<W> iter3(num_qubits, min_weight, max_weight);
+    for (int i = 0; i < 4 * pow(3, 3) + pow(3, 4); i++) {
+        iter3.iter_next();
+    }
+    ASSERT_EQ(iter3.result, PauliString<W>::from_str("+YYYY"));
 })
 
 TEST_EACH_WORD_SIZE_W(pauli_string_iter, find_set_bits, {
@@ -248,6 +246,7 @@ TEST_EACH_WORD_SIZE_W(pauli_string_iter, iter_pauli_string_multi_word, {
     while (iter.iter_next()) {
         count++;
     }
+    ASSERT_EQ(count, 216);
     num_qubits = 65;
     min_weight = 3;
     max_weight = 3;

diff --git a/src/stim/stabilizers/pauli_string_pybind_test.py b/src/stim/stabilizers/pauli_string_pybind_test.py
@@ -903,14 +903,14 @@ def to_py_int(bool_arr):
 
 @pytest.mark.parametrize(
     "num_qubits,min_weight,max_weight",
-    #[(4, 3, 4), (8, 0, 8), (72, 1, 3), (129, 1, 2), (1023, 0, 1), (128, 1, 1)],
-    [(72, 1, 1)]
+    [(4, 3, 4), (8, 0, 8), (72, 1, 3), (129, 1, 2), (1023, 0, 1), (128, 1, 1)],
+    # [(72, 1, 1)]
 )
 def test_iter_all(num_qubits, min_weight, max_weight):
     terms = ["X", "Z", "Y"]
-    pauli_it = stim.PauliString.iter_all(
+    pauli_it = list(stim.PauliString.iter_all(
         num_qubits, min_weight=min_weight, max_weight=max_weight
-    )
+    ))
     pauli_py = []
     for weight in range(min_weight, max_weight + 1):
         # Find all possible locations for the 3^weight pauli products
@@ -921,25 +921,10 @@ def test_iter_all(num_qubits, min_weight, max_weight):
                 tern_rep = np.base_repr(i, 3)
                 pauli = [f"{terms[int(it)]}{qb}" for it, qb in zip(f"{tern_rep:0>{weight}}", comb)]
                 pauli_py.append(pauli)
-
-    # def to_py_int(bool_arr):
-    #     return int("".join(f"{int(b)}" for b in bool_arr[::-1]), 2)
-    # def next_perm(v):
-    #     t = (v | (v - 1)) + 1
-    #     w = t | ((((t & -t) // (v & -v)) >> 1) - 1)
-    #     return w
-    # w = 1
-    # for i in range(64):
-    #     v = next_perm(w)
-    #     print(i, v, bin(v))
-    #     w = v
-        # print(to_py_int(pauli_it.next_qubit_permutation()))
-    # assert len(pauli_it) == len(pauli_py)
-    for i in pauli_it:
-        print(i)
-    # for i, (stim_p, py_p) in enumerate(zip(pauli_it, pauli_py)):
-    #     stim_string = ["_"] * num_qubits
-    #     for pauli in py_p[::-1]:
-    #         pauli_type, indx = pauli[0], int(pauli[1:])
-    #         stim_string[indx] = pauli_type
-    #     assert stim_p == stim.PauliString("".join(s for s in stim_string))
+    assert len(pauli_it) == len(pauli_py)
+    for i, (stim_p, py_p) in enumerate(zip(pauli_it, pauli_py)):
+        stim_string = ["_"] * num_qubits
+        for pauli in py_p[::-1]:
+            pauli_type, indx = pauli[0], int(pauli[1:])
+            stim_string[indx] = pauli_type
+        assert stim_p == stim.PauliString("".join(s for s in stim_string))