Improve performance of differentiable matrix calculations for Pauli Sentences

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

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

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 @@
         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 @@
         """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 @@
             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 @@
         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,36 @@
         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)
+        elif ml_interface == "tf":
+            data0 = qml.math.cast_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)
+            elif ml_interface == "tf":
+                csr_data = qml.math.cast_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(

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)