Feature/pauli string python

benchmarks/pauli_helpers.py

@@ -0,0 +1,17 @@
+import numpy as np
+
+
+def pauli_matrices() -> dict:
+    s0 = np.array([[1, 0], [0, 1]], dtype=np.complex128)
+    s1 = np.array([[0, 1], [1, 0]], dtype=np.complex128)
+    s2 = np.array([[0, -1j], [1j, 0]], dtype=np.complex128)
+    s3 = np.array([[1, 0], [0, -1]], dtype=np.complex128)
+    return {"I": s0, "X": s1, "Y": s2, "Z": s3, 0: s0, 1: s1, 2: s2, 3: s3}
+
+
+def naive_pauli_converter(string: str) -> np.ndarray:
+    paulis = pauli_matrices()
+    matrix = paulis[string[-1]]
+    for p in reversed(string[:-1]):
+        matrix = np.kron(paulis[p], matrix)
+    return matrix

benchmarks/pauli_operations.py

@@ -0,0 +1,103 @@
+import numpy as np
+
+
+class PauliString:
+    def __init__(self, string: str) -> None:
+        if not all([c in "IXYZ" for c in string]):
+            raise ValueError(f"Invalid pauli string {string}")
+
+        self.string = string
+        self.dim = 1 << len(string)
+        self.weight = len(string) - string.count("I")
+
+    def dense(self) -> np.ndarray:
+        columns, values = compose_sparse_pauli(self.string)
+
+        matrix = np.zeros((columns.size, values.size), dtype=np.complex128)
+        matrix[np.arange(columns.size), columns] = values
+        return matrix
+
+    def multiply(self, state: np.ndarray) -> np.ndarray:
+        if state.shape[0] != self.dim or state.ndim > 2:
+            raise ValueError(f"Provided state has inconsistent shape {state.shape}")
+
+        columns, values = compose_sparse_pauli(self.string)
+
+        if state.ndim == 2:
+            return values[:, np.newaxis] * state[columns]
+        else:
+            return values * state[columns]
+
+
+def compose_sparse_pauli(string: str) -> tuple[np.ndarray, np.ndarray]:
+    n_qubits = len(string)
+    n_vals = 1 << n_qubits
+    n_ys = string.count("Y")
+
+    # initialize cols array with zeros as we need first element to be 0
+    cols = np.zeros(n_vals, dtype=np.int32)
+    vals = np.empty(n_vals, dtype=np.complex128)
+
+    for p in string:
+        cols[0] <<= 1
+        if p == "X" or p == "Y":
+            cols[0] += 1
+
+    match n_ys % 4:
+        case 0:
+            vals[0] = 1.0
+        case 1:
+            vals[0] = -1.0j
+        case 2:
+            vals[0] = -1.0
+        case 3:
+            vals[0] = 1.0j
+
+    for q in range(n_qubits):
+        p = string[n_qubits - q - 1]
+        pow_of_two = 1 << q
+
+        new_slice = slice(pow_of_two, 2 * pow_of_two)
+        old_slice = slice(0, pow_of_two)
+
+        match p:
+            case "I":
+                cols[new_slice] = cols[old_slice] + pow_of_two
+                vals[new_slice] = vals[old_slice]
+            case "X":
+                cols[new_slice] = cols[old_slice] - pow_of_two
+                vals[new_slice] = vals[old_slice]
+            case "Y":
+                cols[new_slice] = cols[old_slice] - pow_of_two
+                vals[new_slice] = -vals[old_slice]
+            case "Z":
+                cols[new_slice] = cols[old_slice] + pow_of_two
+                vals[new_slice] = -vals[old_slice]
+
+    return cols, vals
+
+
+def compose_sparse_diag_pauli(string) -> np.ndarray:
+    if "X" in string or "Y" in string:
+        raise ValueError("Pauli string must contain only I and Z characters")
+
+    n_qubits = len(string)
+    n_vals = 1 << n_qubits
+
+    # initialize vals array with ones as we need first element to be 1
+    vals = np.ones(n_vals, dtype=np.complex128)
+
+    for q in range(n_qubits):
+        p = string[n_qubits - q - 1]
+        pow_of_two = 1 << q
+
+        new_slice = slice(pow_of_two, 2 * pow_of_two)
+        old_slice = slice(0, pow_of_two)
+
+        match p:
+            case "I":
+                vals[new_slice] = vals[old_slice]
+            case "Z":
+                vals[new_slice] = -vals[old_slice]
+
+    return vals

benchmarks/test_pauli_operations.py

@@ -0,0 +1,159 @@
+import pytest
+import numpy as np
+from itertools import permutations, chain
+
+from pauli_operations import (
+    PauliString,
+    compose_sparse_pauli,
+    compose_sparse_diag_pauli,
+)
+from pauli_helpers import pauli_matrices, naive_pauli_converter
+
+
+@pytest.fixture
+def paulis():
+    return pauli_matrices()
+
+
+def test_pauli_string(paulis):
+    for p in ["I", "X", "Y", "Z"]:
+        ps = PauliString(p)
+        assert ps.dim == 2
+        assert ps.weight == 1 or p == "I"
+        np.testing.assert_array_equal(ps.dense(), paulis[p])
+        np.testing.assert_array_equal(naive_pauli_converter(p), paulis[p])
+
+    ps = PauliString("III")
+    assert ps.dim == 8
+    assert ps.weight == 0
+    np.testing.assert_array_equal(ps.dense(), np.eye(8))
+    np.testing.assert_array_equal(naive_pauli_converter(ps.string), np.eye(8))
+
+    ps = PauliString(string="IZ")
+    assert ps.dim == 4
+    assert ps.weight == 1
+    np.testing.assert_array_equal(ps.dense(), np.kron(paulis["I"], paulis["Z"]))
+    np.testing.assert_array_equal(
+        naive_pauli_converter(ps.string), np.kron(paulis["I"], paulis["Z"])
+    )
+
+    ps = PauliString(string="XYZ")
+    assert ps.dim == 8
+    assert ps.weight == 3
+    np.testing.assert_array_equal(
+        ps.dense(), np.kron(paulis["X"], np.kron(paulis["Y"], paulis["Z"]))
+    )
+    np.testing.assert_array_equal(
+        naive_pauli_converter(ps.string),
+        np.kron(paulis["X"], np.kron(paulis["Y"], paulis["Z"])),
+    )
+
+    assert PauliString("XYIZXYZ").dim == 2**7
+    assert PauliString("XYIZXYZ").weight == 6
+    assert PauliString("XXIYYIZZ").dim == 2**8
+    assert PauliString("XXIYYIZZ").weight == 6
+    assert PauliString("ZIXIZYXXY").dim == 2**9
+    assert PauliString("ZIXIZYXXY").weight == 7
+
+
+def test_sparse_pauli_composer(paulis):
+    ps = PauliString("II")
+    assert ps.dim == 4
+    np.testing.assert_array_equal(ps.dense(), np.eye(4))
+    np.testing.assert_array_equal(compose_sparse_diag_pauli(ps.string), np.ones(4))
+
+    ps = PauliString("IIII")
+    assert ps.dim == 16
+    np.testing.assert_array_equal(ps.dense(), np.eye(16))
+    np.testing.assert_array_equal(compose_sparse_diag_pauli(ps.string), np.ones(16))
+
+    ps = PauliString("XXX")
+    assert ps.dim == 8
+    np.testing.assert_array_equal(ps.dense(), np.fliplr(np.eye(8)))
+
+    ps = PauliString("IY")
+    np.testing.assert_array_equal(
+        ps.dense(),
+        np.block([[paulis["Y"], np.zeros((2, 2))], [np.zeros((2, 2)), paulis["Y"]]]),
+    )
+
+    ps = PauliString("IZ")
+    np.testing.assert_array_equal(
+        ps.dense(),
+        np.block([[paulis["Z"], np.zeros((2, 2))], [np.zeros((2, 2)), paulis["Z"]]]),
+    )
+
+    np.testing.assert_array_equal(
+        compose_sparse_diag_pauli("IZI"),
+        compose_sparse_pauli("IZI")[1],
+    )
+
+    np.testing.assert_array_equal(
+        compose_sparse_diag_pauli("ZIZI"),
+        compose_sparse_pauli("ZIZI")[1],
+    )
+
+    np.testing.assert_array_equal(
+        compose_sparse_diag_pauli("ZZZIII"),
+        compose_sparse_pauli("ZZZIII")[1],
+    )
+
+
+def test_sparse_pauli_composer_equivalence():
+    for c in ["I", "X", "Y", "Z"]:
+        np.testing.assert_array_equal(
+            PauliString(c).dense(),
+            naive_pauli_converter(c),
+        )
+
+    for s in permutations("XYZ", 2):
+        s = "".join(s)
+        np.testing.assert_array_equal(
+            PauliString(s).dense(),
+            naive_pauli_converter(s),
+        )
+
+    for s in permutations("IXYZ", 3):
+        s = "".join(s)
+        np.testing.assert_array_equal(
+            PauliString(s).dense(),
+            naive_pauli_converter(s),
+        )
+
+    ixyz = PauliString("IXYZ").dense()
+    np.testing.assert_array_equal(ixyz, naive_pauli_converter("IXYZ"))
+
+    zyxi = PauliString("ZYXI").dense()
+    np.testing.assert_array_equal(zyxi, naive_pauli_converter("ZYXI"))
+
+    assert np.abs(ixyz - zyxi).sum().sum() > 1e-10
+
+    for s in ["XYIZXYZ", "XXIYYIZZ", "ZIXIZYXX"]:
+        np.testing.assert_array_equal(PauliString(s).dense(), naive_pauli_converter(s))
+
+
+def test_sparse_pauli_multiply():
+    rng = np.random.default_rng(321)
+
+    for s in chain(
+        list("IXYZ"), list(permutations("IXYZ", 3)), ["XYIZXYZ", "XXIYYIZZ", "ZIXIZYXX"]
+    ):
+        s = "".join(s)
+        n = 2 ** len(s)
+        psi = rng.random(n)
+        psi_batch = rng.random((n, 21))
+
+        np.testing.assert_allclose(
+            PauliString(s).multiply(psi),
+            naive_pauli_converter(s).dot(psi),
+            atol=1e-15,
+        )
+        np.testing.assert_allclose(
+            PauliString(s).multiply(psi_batch),
+            naive_pauli_converter(s) @ psi_batch,
+            atol=1e-15,
+        )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])