diff --git a/benchmarks/pauli_helpers.py b/benchmarks/pauli_helpers.py new file mode 100644 index 0000000..27cf0e4 --- /dev/null +++ b/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 diff --git a/benchmarks/pauli_operations.py b/benchmarks/pauli_operations.py new file mode 100644 index 0000000..c9db78f --- /dev/null +++ b/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 diff --git a/benchmarks/test_pauli_operations.py b/benchmarks/test_pauli_operations.py new file mode 100644 index 0000000..df7f300 --- /dev/null +++ b/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__])