Add stim.PauliString.pauli_indices (#710)

Fixes #699
quantumlib · Mar 12, 2024 · b29741d · b29741d
1 parent 7c90532
commit b29741d
Show file tree

Hide file tree

Showing 5 changed files with 245 additions and 0 deletions.
diff --git a/doc/python_api_reference_vDev.md b/doc/python_api_reference_vDev.md
@@ -264,6 +264,7 @@ API references for stable versions are kept on the [stim github wiki](https://gi
     - [`stim.PauliString.from_numpy`](#stim.PauliString.from_numpy)
     - [`stim.PauliString.from_unitary_matrix`](#stim.PauliString.from_unitary_matrix)
     - [`stim.PauliString.iter_all`](#stim.PauliString.iter_all)
+    - [`stim.PauliString.pauli_indices`](#stim.PauliString.pauli_indices)
     - [`stim.PauliString.random`](#stim.PauliString.random)
     - [`stim.PauliString.sign`](#stim.PauliString.sign)
     - [`stim.PauliString.to_numpy`](#stim.PauliString.to_numpy)
@@ -9068,6 +9069,50 @@ def iter_all(
     """
 ```
 
+<a name="stim.PauliString.pauli_indices"></a>
+```python
+# stim.PauliString.pauli_indices
+
+# (in class stim.PauliString)
+def pauli_indices(
+    self,
+    included_paulis: str = "XYZ",
+) -> List[int]:
+    """Returns the indices of non-identity Paulis, or of specified Paulis.
+
+    Args:
+        include: A string containing the Pauli types to include.
+            X type Pauli indices are included if "X" or "x" is in the string.
+            Y type Pauli indices are included if "Y" or "y" is in the string.
+            Z type Pauli indices are included if "Z" or "z" is in the string.
+            I type Pauli indices are included if "I" or "_" is in the string.
+            An exception is thrown if other characters are in the string.
+
+    Returns:
+        A list containing the ascending indices of matching Pauli terms.
+
+    Examples:
+        >>> import stim
+        >>> stim.PauliString("_____X___Y____Z___").pauli_indices()
+        [5, 9, 14]
+
+        >>> stim.PauliString("_____X___Y____Z___").pauli_indices("XZ")
+        [5, 14]
+
+        >>> stim.PauliString("_____X___Y____Z___").pauli_indices("X")
+        [5]
+
+        >>> stim.PauliString("_____X___Y____Z___").pauli_indices("Y")
+        [9]
+
+        >>> stim.PauliString("_____X___Y____Z___").pauli_indices("IY")
+        [0, 1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 15, 16, 17]
+
+        >>> stim.PauliString("-X103*Y100").pauli_indices()
+        [100, 103]
+    """
+```
+
 <a name="stim.PauliString.random"></a>
 ```python
 # stim.PauliString.random

diff --git a/doc/stim.pyi b/doc/stim.pyi
@@ -6982,6 +6982,43 @@ class PauliString:
             +_ZX
             +_ZZ
         """
+    def pauli_indices(
+        self,
+        included_paulis: str = "XYZ",
+    ) -> List[int]:
+        """Returns the indices of non-identity Paulis, or of specified Paulis.
+
+        Args:
+            include: A string containing the Pauli types to include.
+                X type Pauli indices are included if "X" or "x" is in the string.
+                Y type Pauli indices are included if "Y" or "y" is in the string.
+                Z type Pauli indices are included if "Z" or "z" is in the string.
+                I type Pauli indices are included if "I" or "_" is in the string.
+                An exception is thrown if other characters are in the string.
+
+        Returns:
+            A list containing the ascending indices of matching Pauli terms.
+
+        Examples:
+            >>> import stim
+            >>> stim.PauliString("_____X___Y____Z___").pauli_indices()
+            [5, 9, 14]
+
+            >>> stim.PauliString("_____X___Y____Z___").pauli_indices("XZ")
+            [5, 14]
+
+            >>> stim.PauliString("_____X___Y____Z___").pauli_indices("X")
+            [5]
+
+            >>> stim.PauliString("_____X___Y____Z___").pauli_indices("Y")
+            [9]
+
+            >>> stim.PauliString("_____X___Y____Z___").pauli_indices("IY")
+            [0, 1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 15, 16, 17]
+
+            >>> stim.PauliString("-X103*Y100").pauli_indices()
+            [100, 103]
+        """
     @staticmethod
     def random(
         num_qubits: int,

diff --git a/glue/python/src/stim/__init__.pyi b/glue/python/src/stim/__init__.pyi
@@ -6982,6 +6982,43 @@ class PauliString:
             +_ZX
             +_ZZ
         """
+    def pauli_indices(
+        self,
+        included_paulis: str = "XYZ",
+    ) -> List[int]:
+        """Returns the indices of non-identity Paulis, or of specified Paulis.
+
+        Args:
+            include: A string containing the Pauli types to include.
+                X type Pauli indices are included if "X" or "x" is in the string.
+                Y type Pauli indices are included if "Y" or "y" is in the string.
+                Z type Pauli indices are included if "Z" or "z" is in the string.
+                I type Pauli indices are included if "I" or "_" is in the string.
+                An exception is thrown if other characters are in the string.
+
+        Returns:
+            A list containing the ascending indices of matching Pauli terms.
+
+        Examples:
+            >>> import stim
+            >>> stim.PauliString("_____X___Y____Z___").pauli_indices()
+            [5, 9, 14]
+
+            >>> stim.PauliString("_____X___Y____Z___").pauli_indices("XZ")
+            [5, 14]
+
+            >>> stim.PauliString("_____X___Y____Z___").pauli_indices("X")
+            [5]
+
+            >>> stim.PauliString("_____X___Y____Z___").pauli_indices("Y")
+            [9]
+
+            >>> stim.PauliString("_____X___Y____Z___").pauli_indices("IY")
+            [0, 1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 15, 16, 17]
+
+            >>> stim.PauliString("-X103*Y100").pauli_indices()
+            [100, 103]
+        """
     @staticmethod
     def random(
         num_qubits: int,

diff --git a/src/stim/stabilizers/pauli_string.pybind.cc b/src/stim/stabilizers/pauli_string.pybind.cc
@@ -636,6 +636,103 @@ void stim_pybind::pybind_pauli_string_methods(pybind11::module &m, pybind11::cla
         )DOC")
             .data());
 
+    c.def(
+        "pauli_indices",
+        [](const FlexPauliString &self, const std::string &include) {
+            std::vector<uint64_t> result;
+            size_t n64 = self.value.xs.num_u64_padded();
+            bool keep_i = false;
+            bool keep_x = false;
+            bool keep_y = false;
+            bool keep_z = false;
+            for (char c : include) {
+                switch (c) {
+                    case '_':
+                    case 'I':
+                        keep_i = true;
+                        break;
+                    case 'x':
+                    case 'X':
+                        keep_x = true;
+                        break;
+                    case 'y':
+                    case 'Y':
+                        keep_y = true;
+                        break;
+                    case 'z':
+                    case 'Z':
+                        keep_z = true;
+                        break;
+                    default:
+                        throw std::invalid_argument("Invalid character in include string: " + std::string(1, c));
+                }
+            }
+            for (size_t k = 0; k < n64; k++) {
+                uint64_t x = self.value.xs.u64[k];
+                uint64_t z = self.value.zs.u64[k];
+                uint64_t u = 0;
+                if (keep_i) {
+                    u |= ~x & ~z;
+                }
+                if (keep_x) {
+                    u |= x & ~z;
+                }
+                if (keep_y) {
+                    u |= x & z;
+                }
+                if (keep_z) {
+                    u |= ~x & z;
+                }
+                while (u) {
+                    uint8_t v = std::countr_zero(u);
+                    uint64_t q = k * 64 + v;
+                    if (q >= self.value.num_qubits) {
+                        return result;
+                    }
+                    result.push_back(q);
+                    u &= u - 1;
+                }
+            }
+            return result;
+        },
+        pybind11::arg("included_paulis") = "XYZ",
+        clean_doc_string(R"DOC(
+            @signature def pauli_indices(self, included_paulis: str = "XYZ") -> List[int]:
+            Returns the indices of non-identity Paulis, or of specified Paulis.
+
+            Args:
+                include: A string containing the Pauli types to include.
+                    X type Pauli indices are included if "X" or "x" is in the string.
+                    Y type Pauli indices are included if "Y" or "y" is in the string.
+                    Z type Pauli indices are included if "Z" or "z" is in the string.
+                    I type Pauli indices are included if "I" or "_" is in the string.
+                    An exception is thrown if other characters are in the string.
+
+            Returns:
+                A list containing the ascending indices of matching Pauli terms.
+
+            Examples:
+                >>> import stim
+                >>> stim.PauliString("_____X___Y____Z___").pauli_indices()
+                [5, 9, 14]
+
+                >>> stim.PauliString("_____X___Y____Z___").pauli_indices("XZ")
+                [5, 14]
+
+                >>> stim.PauliString("_____X___Y____Z___").pauli_indices("X")
+                [5]
+
+                >>> stim.PauliString("_____X___Y____Z___").pauli_indices("Y")
+                [9]
+
+                >>> stim.PauliString("_____X___Y____Z___").pauli_indices("IY")
+                [0, 1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 15, 16, 17]
+
+                >>> stim.PauliString("-X103*Y100").pauli_indices()
+                [100, 103]
+        )DOC")
+            .data());
+
     c.def(
         "commutes",
         [](const FlexPauliString &self, const FlexPauliString &other) {

diff --git a/src/stim/stabilizers/pauli_string_pybind_test.py b/src/stim/stabilizers/pauli_string_pybind_test.py
@@ -899,3 +899,32 @@ def test_backwards_compatibility_init():
     assert stim.PauliString(text="XYZ") == stim.PauliString("+XYZ")
     # noinspection PyArgumentList
     assert stim.PauliString(other=stim.PauliString("XYZ")) == stim.PauliString("+XYZ")
+
+
+def test_pauli_indices():
+    assert stim.PauliString().pauli_indices() == []
+    assert stim.PauliString().pauli_indices("X") == []
+    assert stim.PauliString().pauli_indices("I") == []
+    assert stim.PauliString(5).pauli_indices() == []
+    assert stim.PauliString(5).pauli_indices("X") == []
+    assert stim.PauliString(5).pauli_indices("I") == [0, 1, 2, 3, 4]
+    assert stim.PauliString("X1000").pauli_indices() == [1000]
+    assert stim.PauliString("Y1000").pauli_indices() == [1000]
+    assert stim.PauliString("Z1000").pauli_indices() == [1000]
+    assert stim.PauliString("X1000").pauli_indices("YZ") == []
+    assert stim.PauliString("Y1000").pauli_indices("XZ") == []
+    assert stim.PauliString("Z1000").pauli_indices("XY") == []
+    assert stim.PauliString("X1000").pauli_indices("X") == [1000]
+    assert stim.PauliString("Y1000").pauli_indices("Y") == [1000]
+    assert stim.PauliString("Z1000").pauli_indices("Z") == [1000]
+
+    assert stim.PauliString("_XYZ").pauli_indices("x") == [1]
+    assert stim.PauliString("_XYZ").pauli_indices("X") == [1]
+    assert stim.PauliString("_XYZ").pauli_indices("y") == [2]
+    assert stim.PauliString("_XYZ").pauli_indices("Y") == [2]
+    assert stim.PauliString("_XYZ").pauli_indices("z") == [3]
+    assert stim.PauliString("_XYZ").pauli_indices("Z") == [3]
+    assert stim.PauliString("_XYZ").pauli_indices("I") == [0]
+    assert stim.PauliString("_XYZ").pauli_indices("_") == [0]
+    with pytest.raises(ValueError, match="Invalid character"):
+        assert stim.PauliString("_XYZ").pauli_indices("k")