diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 7d475e15105..ec96fe30621 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -383,6 +383,10 @@

Other improvements

+* Calculating the dense, differentiable matrix for `PauliSentence` and operators with pauli sentences + is now faster. + [(#5578)](https://github.com/PennyLaneAI/pennylane/pull/5578) + * `DefaultQubit` now uses the provided seed for sampling mid-circuit measurements with finite shots. [(#5337)](https://github.com/PennyLaneAI/pennylane/pull/5337) diff --git a/pennylane/ops/op_math/prod.py b/pennylane/ops/op_math/prod.py index 3c9bc0a19a4..df07a8226fd 100644 --- a/pennylane/ops/op_math/prod.py +++ b/pennylane/ops/op_math/prod.py @@ -305,6 +305,8 @@ def decomposition(self): def matrix(self, wire_order=None): """Representation of the operator as a matrix in the computational basis.""" + if self.pauli_rep: + return self.pauli_rep.to_mat(wire_order=wire_order or self.wires) mats: List[TensorLike] = [] batched: List[bool] = [] # batched[i] tells if mats[i] is batched or not diff --git a/pennylane/ops/op_math/sum.py b/pennylane/ops/op_math/sum.py index 19b4ddde1ff..5d3d89dc42b 100644 --- a/pennylane/ops/op_math/sum.py +++ b/pennylane/ops/op_math/sum.py @@ -325,6 +325,8 @@ def matrix(self, wire_order=None): Returns: tensor_like: matrix representation """ + if self.pauli_rep: + return self.pauli_rep.to_mat(wire_order=wire_order or self.wires) gen = ( (qml.matrix(op) if isinstance(op, qml.ops.Hamiltonian) else op.matrix(), op.wires) for op in self diff --git a/pennylane/pauli/pauli_arithmetic.py b/pennylane/pauli/pauli_arithmetic.py index 418550bf3aa..9878b744b37 100644 --- a/pennylane/pauli/pauli_arithmetic.py +++ b/pennylane/pauli/pauli_arithmetic.py @@ -15,7 +15,6 @@ # pylint:disable=protected-access from copy import copy from functools import reduce, lru_cache -from typing import Iterable import numpy as np from scipy import sparse @@ -455,6 +454,8 @@ def _get_csr_data(self, wire_order, coeff): full_word = [self[wire] for wire in wire_order] matrix_size = 2 ** len(wire_order) + if len(self) == 0: + return np.full(matrix_size, coeff, dtype=np.complex128) data = np.empty(matrix_size, dtype=np.complex128) # Non-zero values current_size = 2 data[:current_size], _ = _cached_sparse_data(full_word[-1]) @@ -486,6 +487,8 @@ def _get_csr_indices(self, wire_order): """Computes the sparse matrix indices of the Pauli word times a coefficient, given a wire order.""" full_word = [self[wire] for wire in wire_order] matrix_size = 2 ** len(wire_order) + if len(self) == 0: + return _cached_arange(matrix_size) indices = np.empty(matrix_size, dtype=np.int64) # Column index of non-zero values current_size = 2 _, indices[:current_size] = _cached_sparse_data(full_word[-1]) @@ -838,51 +841,15 @@ def to_mat(self, wire_order=None, format="dense", buffer_size=None): ValueError: Can't get the matrix of an empty PauliSentence. """ wire_order = self.wires if wire_order is None else Wires(wire_order) - if not wire_order.contains_wires(self.wires): - raise ValueError( - "Can't get the matrix for the specified wire order because it " - f"does not contain all the Pauli sentence's wires {self.wires}" - ) - - def _pw_wires(w: Iterable) -> Wires: - """Return the native Wires instance for a list of wire labels. - w represents the wires of the PauliWord being processed. In case - the PauliWord is empty ({}), choose any arbitrary wire from the - PauliSentence it is composed in. - """ - if w: - # PauliWord is not empty, so we can use its wires - return Wires(w) - - if wire_order: - # PauliWord is empty, treat it as Identity operator on any wire - # Pick any arbitrary wire from wire_order - return Wires(wire_order[0]) - - return wire_order - if len(self) == 0: n = len(wire_order) if wire_order is not None else 0 if format == "dense": return np.zeros((2**n, 2**n)) return sparse.csr_matrix((2**n, 2**n), dtype="complex128") - if format != "dense": - return self._to_sparse_mat(wire_order, buffer_size=buffer_size) - - mats_and_wires_gen = ( - ( - coeff * pw.to_mat(wire_order=_pw_wires(pw.wires), format=format), - _pw_wires(pw.wires), - ) - for pw, coeff in self.items() - ) - - reduced_mat, result_wire_order = math.reduce_matrices( - mats_and_wires_gen=mats_and_wires_gen, reduce_func=math.add - ) - - return math.expand_matrix(reduced_mat, result_wire_order, wire_order=wire_order) + if format == "dense": + return self._to_dense_mat(wire_order) + return self._to_sparse_mat(wire_order, buffer_size=buffer_size) def _to_sparse_mat(self, wire_order, buffer_size=None): """Compute the sparse matrix of the Pauli sentence by efficiently adding the Pauli words @@ -922,6 +889,31 @@ def _to_sparse_mat(self, wire_order, buffer_size=None): matrix.eliminate_zeros() return matrix + def _to_dense_mat(self, wire_order): + """Compute the dense matrix of the Pauli sentence by efficiently adding the Pauli words + that it is composed of. See pauli_sparse_matrices.md for the technical details.""" + pauli_words = list(self) # Ensure consistent ordering + + try: + op_sparse_idx = _ps_to_sparse_index(pauli_words, wire_order) + except qml.wires.WireError as e: + raise ValueError( + "Can't get the matrix for the specified wire order because it " + f"does not contain all the Pauli sentence's wires {self.wires}" + ) from e + _, unique_sparse_structures, unique_invs = np.unique( + op_sparse_idx, axis=0, return_index=True, return_inverse=True + ) + pw_sparse_structures = unique_sparse_structures[unique_invs] + + full_matrix = None + for sparse_structure in unique_sparse_structures: + indices, *_ = np.nonzero(pw_sparse_structures == sparse_structure) + mat = self._sum_same_structure_pws_dense([pauli_words[i] for i in indices], wire_order) + + full_matrix = mat if full_matrix is None else qml.math.add(full_matrix, mat) + return full_matrix + def dot(self, vector, wire_order=None): """Computes the matrix-vector product of the Pauli sentence with a state vector. See pauli_sparse_matrices.md for the technical details.""" @@ -964,6 +956,32 @@ def _get_same_structure_csr(self, pauli_words, wire_order): data = outer.T @ inner return indices, data.ravel() + def _sum_same_structure_pws_dense(self, pauli_words, wire_order): + matrix_size = 2 ** (len(wire_order)) + base_matrix = sparse.csr_matrix((matrix_size, matrix_size), dtype="complex128") + + data0 = pauli_words[0]._get_csr_data(wire_order, 1) + base_matrix.data = np.ones_like(data0) + base_matrix.indices = pauli_words[0]._get_csr_indices(wire_order) + base_matrix.indptr = _cached_arange( + matrix_size + 1 + ) # Non-zero entries by row (starting from 0) + base_matrix = base_matrix.toarray() + coeff = self[pauli_words[0]] + ml_interface = qml.math.get_interface(coeff) + if ml_interface == "torch": + data0 = qml.math.convert_like(data0, coeff) + data = coeff * data0 + for pw in pauli_words[1:]: + coeff = self[pw] + csr_data = pw._get_csr_data(wire_order, 1) + ml_interface = qml.math.get_interface(coeff) + if ml_interface == "torch": + csr_data = qml.math.convert_like(csr_data, coeff) + data += self[pw] * csr_data + + return qml.math.einsum("ij,i->ij", base_matrix, data) + def _sum_same_structure_pws(self, pauli_words, wire_order): """Sums Pauli words with the same sparse structure.""" mat = pauli_words[0].to_mat( diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py index c411c03ba7a..2f73e465e6e 100644 --- a/tests/devices/qubit/test_apply_operation.py +++ b/tests/devices/qubit/test_apply_operation.py @@ -389,7 +389,8 @@ def test_large_state_small_matrix_evolves_matrix(self, mocker): # seems like _evolve_state_vector_under_parametrized_evolution calls # einsum twice, and the default apply_operation only once - assert spy.call_count == 1 + # and it seems that getting the matrix from the hamiltonian calls einsum a few times. + assert spy.call_count == 6 def test_small_evolves_state(self, mocker): """Test that applying a ParametrizedEvolution operating on less @@ -465,7 +466,8 @@ def test_small_evolves_state(self, mocker): # seems like _evolve_state_vector_under_parametrized_evolution calls # einsum twice, and the default apply_operation only once - assert spy.call_count == 2 + # and it seems that getting the matrix from the hamiltonian calls einsum a few times. + assert spy.call_count == 7 def test_parametrized_evolution_raises_error(self): """Test applying a ParametrizedEvolution without params or t specified raises an error.""" @@ -530,9 +532,11 @@ def test_with_batched_state(self, num_state_wires, mocker): assert np.allclose(new_state, new_state_expected, atol=0.002) if num_state_wires == 4: - assert spy_einsum.call_count == 2 + # and it seems that getting the matrix from the hamiltonian calls einsum a few times. + assert spy_einsum.call_count == 7 else: - assert spy_einsum.call_count == 1 + # and it seems that getting the matrix from the hamiltonian calls einsum a few times. + assert spy_einsum.call_count == 6 @pytest.mark.parametrize("ml_framework", ml_frameworks_list) diff --git a/tests/pauli/test_pauli_arithmetic.py b/tests/pauli/test_pauli_arithmetic.py index 67d0c40db98..62229838b0c 100644 --- a/tests/pauli/test_pauli_arithmetic.py +++ b/tests/pauli/test_pauli_arithmetic.py @@ -901,93 +901,6 @@ def test_iadd_ps_pw(self, ps, pw, res): assert copy_ps1 == res # Check if the modified object matches the expected result assert copy_ps2 == ps # Ensure the original object is not modified - PS_EMPTY_CASES = ( - (PauliSentence({}), np.zeros((1, 1))), - (PauliSentence({PauliWord({}): 1.0}), np.ones((1, 1))), - (PauliSentence({PauliWord({}): 2.5}), 2.5 * np.ones((1, 1))), - ( - PauliSentence({PauliWord({}): 2.5, PauliWord({0: "X"}): 1.0}), - 2.5 * np.eye(2) + qml.matrix(qml.PauliX(0)), - ), - ) - - @pytest.mark.parametrize("ps, true_res", PS_EMPTY_CASES) - def test_to_mat_empty(self, ps, true_res): - """Test that empty PauliSentences and PauliSentences with empty PauliWords are handled correctly""" - - res_dense = ps.to_mat() - assert np.allclose(res_dense, true_res) - res_sparse = ps.to_mat(format="csr") - assert sparse.issparse(res_sparse) - assert qml.math.allclose(res_sparse.todense(), true_res) - - def test_empty_pauli_to_mat_with_wire_order(self): - """Test the to_mat method with an empty PauliSentence and PauliWord and an external wire order.""" - actual = PauliSentence({PauliWord({}): 1.5}).to_mat([0, 1]) - assert np.allclose(actual, 1.5 * np.eye(4)) - - ps_wire_order = ((ps1, []), (ps1, [0, 1, 2, "a", "b"]), (ps3, [0, 1, "c"])) - - @pytest.mark.parametrize("ps, wire_order", ps_wire_order) - def test_to_mat_error_incomplete(self, ps, wire_order): - """Test that an appropriate error is raised when the wire order does - not contain all the PauliSentence's wires.""" - match = "Can't get the matrix for the specified wire order" - with pytest.raises(ValueError, match=match): - ps.to_mat(wire_order=wire_order) - - def test_to_mat_empty_sentence_with_wires(self): - """Test that a zero matrix is returned if wire_order is provided on an empty PauliSentence.""" - true_res = np.zeros((4, 4)) - res_dense = ps5.to_mat(wire_order=[0, 1]) - assert np.allclose(res_dense, true_res) - res_sparse = ps5.to_mat(wire_order=[0, 1], format="csr") - assert sparse.issparse(res_sparse) - assert qml.math.allclose(res_sparse.todense(), true_res) - - tup_ps_mat = ( - ( - ps1, - [0, 1, 2, "a", "b", "c"], - 1.23 * np.kron(np.kron(matI, np.kron(matX, matY)), np.eye(8)) - + 4j * np.kron(np.eye(8), np.kron(matX, np.kron(matX, matZ))) - - 0.5 * np.kron(matZ, np.kron(np.eye(8), np.kron(matZ, matZ))), - ), - ( - ps2, - ["a", "b", "c", 0, 1, 2], - -1.23 * np.kron(np.eye(8), np.kron(matI, np.kron(matX, matY))) - - 4j * np.kron(np.kron(matX, np.kron(matX, matZ)), np.eye(8)) - + 0.5 * np.kron(np.kron(matI, np.kron(matZ, np.kron(matZ, matZ))), np.eye(4)), - ), - ( - ps3, - [0, "b", "c"], - -0.5 * np.kron(matZ, np.kron(matZ, matZ)) + 1 * np.eye(8), - ), - ) - - @pytest.mark.parametrize("ps, wire_order, true_matrix", tup_ps_mat) - def test_to_mat_wire_order(self, ps, wire_order, true_matrix): - """Test that the wire_order is correctly incorporated in computing the - matrix representation.""" - assert np.allclose(ps.to_mat(wire_order), true_matrix) - - @pytest.mark.parametrize("ps, wire_order, true_matrix", tup_ps_mat) - def test_to_mat_format(self, ps, wire_order, true_matrix): - """Test that the correct type of matrix is returned given the format kwarg.""" - sparse_mat = ps.to_mat(wire_order, format="csr") - assert sparse.issparse(sparse_mat) - assert np.allclose(sparse_mat.toarray(), true_matrix) - - @pytest.mark.parametrize("ps,wire_order,true_matrix", tup_ps_mat) - def test_to_mat_buffer(self, ps, wire_order, true_matrix): - """Test that the intermediate matrices are added correctly once the maximum buffer - size is reached.""" - buffer_size = 2 ** len(wire_order) * 48 # Buffer size for 2 matrices - sparse_mat = ps.to_mat(wire_order, format="csr", buffer_size=buffer_size) - assert np.allclose(sparse_mat.toarray(), true_matrix) - def test_simplify(self): """Test that simplify removes terms in the PauliSentence with coefficient less than the threshold""" @@ -1208,6 +1121,171 @@ def test_map_wires(self): ) +class TestPauliSentenceMatrix: + """Tests for calculating the matrix of a pauli sentence.""" + + PS_EMPTY_CASES = ( + (PauliSentence({}), np.zeros((1, 1))), + (PauliSentence({PauliWord({}): 1.0}), np.ones((1, 1))), + (PauliSentence({PauliWord({}): 2.5}), 2.5 * np.ones((1, 1))), + ( + PauliSentence({PauliWord({}): 2.5, PauliWord({0: "X"}): 1.0}), + 2.5 * np.eye(2) + qml.matrix(qml.PauliX(0)), + ), + ) + + @pytest.mark.parametrize("ps, true_res", PS_EMPTY_CASES) + def test_to_mat_empty(self, ps, true_res): + """Test that empty PauliSentences and PauliSentences with empty PauliWords are handled correctly""" + + res_dense = ps.to_mat() + assert np.allclose(res_dense, true_res) + res_sparse = ps.to_mat(format="csr") + assert sparse.issparse(res_sparse) + assert qml.math.allclose(res_sparse.todense(), true_res) + + def test_empty_pauli_to_mat_with_wire_order(self): + """Test the to_mat method with an empty PauliSentence and PauliWord and an external wire order.""" + actual = PauliSentence({PauliWord({}): 1.5}).to_mat([0, 1]) + assert np.allclose(actual, 1.5 * np.eye(4)) + + ps_wire_order = ((ps1, []), (ps1, [0, 1, 2, "a", "b"]), (ps3, [0, 1, "c"])) + + @pytest.mark.parametrize("ps, wire_order", ps_wire_order) + def test_to_mat_error_incomplete(self, ps, wire_order): + """Test that an appropriate error is raised when the wire order does + not contain all the PauliSentence's wires.""" + match = "Can't get the matrix for the specified wire order" + with pytest.raises(ValueError, match=match): + ps.to_mat(wire_order=wire_order) + + def test_to_mat_empty_sentence_with_wires(self): + """Test that a zero matrix is returned if wire_order is provided on an empty PauliSentence.""" + true_res = np.zeros((4, 4)) + res_dense = ps5.to_mat(wire_order=[0, 1]) + assert np.allclose(res_dense, true_res) + res_sparse = ps5.to_mat(wire_order=[0, 1], format="csr") + assert sparse.issparse(res_sparse) + assert qml.math.allclose(res_sparse.todense(), true_res) + + tup_ps_mat = ( + ( + ps1, + [0, 1, 2, "a", "b", "c"], + 1.23 * np.kron(np.kron(matI, np.kron(matX, matY)), np.eye(8)) + + 4j * np.kron(np.eye(8), np.kron(matX, np.kron(matX, matZ))) + - 0.5 * np.kron(matZ, np.kron(np.eye(8), np.kron(matZ, matZ))), + ), + ( + ps2, + ["a", "b", "c", 0, 1, 2], + -1.23 * np.kron(np.eye(8), np.kron(matI, np.kron(matX, matY))) + - 4j * np.kron(np.kron(matX, np.kron(matX, matZ)), np.eye(8)) + + 0.5 * np.kron(np.kron(matI, np.kron(matZ, np.kron(matZ, matZ))), np.eye(4)), + ), + ( + ps3, + [0, "b", "c"], + -0.5 * np.kron(matZ, np.kron(matZ, matZ)) + 1 * np.eye(8), + ), + ) + + @pytest.mark.parametrize("ps, wire_order, true_matrix", tup_ps_mat) + def test_to_mat_wire_order(self, ps, wire_order, true_matrix): + """Test that the wire_order is correctly incorporated in computing the + matrix representation.""" + assert np.allclose(ps.to_mat(wire_order), true_matrix) + + @pytest.mark.parametrize("ps, wire_order, true_matrix", tup_ps_mat) + def test_to_mat_format(self, ps, wire_order, true_matrix): + """Test that the correct type of matrix is returned given the format kwarg.""" + sparse_mat = ps.to_mat(wire_order, format="csr") + assert sparse.issparse(sparse_mat) + assert np.allclose(sparse_mat.toarray(), true_matrix) + + @pytest.mark.parametrize("ps,wire_order,true_matrix", tup_ps_mat) + def test_to_mat_buffer(self, ps, wire_order, true_matrix): + """Test that the intermediate matrices are added correctly once the maximum buffer + size is reached.""" + buffer_size = 2 ** len(wire_order) * 48 # Buffer size for 2 matrices + sparse_mat = ps.to_mat(wire_order, format="csr", buffer_size=buffer_size) + assert np.allclose(sparse_mat.toarray(), true_matrix) + + @pytest.mark.tf + def test_dense_matrix_tf(self): + """Test calculating the matrix for a pauli sentence is differentaible with tensorflow.""" + import tensorflow as tf + + x = tf.Variable(0.1 + 0j) + y = tf.Variable(0.2 + 0j) + + with tf.GradientTape() as tape: + _pw1 = PauliWord({0: "X", 1: "Y"}) + _pw2 = PauliWord({0: "Y", 1: "X"}) + H = x * _pw1 + y * _pw2 + mat = H.to_mat() + + gx, gy = tape.jacobian(mat, [x, y]) + + pw1_mat = np.array([[0, 0, 0, -1j], [0, 0, 1j, 0], [0, -1j, 0, 0], [1j, 0, 0, 0]]) + pw2_mat = np.array([[0, 0, 0, -1j], [0, 0, -1j, 0], [0, 1j, 0, 0], [1j, 0, 0, 0]]) + + assert qml.math.allclose(mat, x * pw1_mat + y * pw2_mat) + assert qml.math.allclose(gx, qml.math.conj(pw1_mat)) # tf complex number convention + assert qml.math.allclose(gy, qml.math.conj(pw2_mat)) + + @pytest.mark.torch + def test_dense_matrix_torch(self): + """Test calculating and differentiating the matrix with torch.""" + + import torch + + _pw1 = qml.pauli.PauliWord({0: "X", 1: "Z"}) + _pw2 = qml.pauli.PauliWord({0: "X"}) + + pw1_mat = torch.tensor([[0, 0, 1, 0], [0, 0, 0, -1], [1, 0, 0, 0], [0, -1, 0, 0]]) + pw2_mat = torch.tensor([[0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0], [0, 1, 0, 0]]) + + x = torch.tensor(0.1, requires_grad=True) + y = torch.tensor(0.2, requires_grad=True) + + def f(x, y): + H = x * _pw1 + y * _pw2 + return qml.math.real(H.to_mat()) + + mat = f(x, y) + assert qml.math.allclose(mat, x * pw1_mat + y * pw2_mat) + + gx, gy = torch.autograd.functional.jacobian(f, (x, y)) + assert qml.math.allclose(gx, pw1_mat) + assert qml.math.allclose(gy, pw2_mat) + + @pytest.mark.jax + def test_dense_matrix_jax(self): + """Test calculating and differentiating the matrix with jax.""" + + import jax + + def f(x, y): + _pw1 = qml.pauli.PauliWord({0: "X", 1: "Y"}) + _pw2 = qml.pauli.PauliWord({0: "Y", 1: "X"}) + H = x * _pw1 + y * _pw2 + return H.to_mat() + + x = jax.numpy.array(0.1 + 0j) + y = jax.numpy.array(0.2 + 0j) + + pw1_mat = np.array([[0, 0, 0, -1j], [0, 0, 1j, 0], [0, -1j, 0, 0], [1j, 0, 0, 0]]) + pw2_mat = np.array([[0, 0, 0, -1j], [0, 0, -1j, 0], [0, 1j, 0, 0], [1j, 0, 0, 0]]) + + mat = f(x, y) + assert qml.math.allclose(mat, x * pw1_mat + y * pw2_mat) + + gx, gy = jax.jacobian(f, holomorphic=True, argnums=(0, 1))(x, y) + assert qml.math.allclose(gx, pw1_mat) + assert qml.math.allclose(gy, pw2_mat) + + class TestPaulicomms: """Test 'native' commutator in PauliWord and PauliSentence"""