[Op. Arithm.] Improve performance #1 (#3022)

* add hash

* use hash to compute eigendecomposition

* use eigvals of factors if they commute

* improve prod

* merge

* add hash to sum

* add tests

* revert

* docs

* refactor

* remove adjoint

* chore (changelog): ✏️ Add feature to changelog.

* chore (changelog): ✏️ Add feature to changelog.

* Update doc/releases/changelog-dev.md

Co-authored-by: Christina Lee <[email protected]>

Co-authored-by: Christina Lee <[email protected]>

Author: AlbertMitjans

doc/releases/changelog-dev.md

@@ -266,7 +266,7 @@
 * `Controlled` operators now work with `qml.is_commuting`.
   [(#2994)](https://github.com/PennyLaneAI/pennylane/pull/2994)
 
-* `Prod` and `Sum` class now support the `sparse_matrix()` method. 
+* `Prod` and `Sum` class now support the `sparse_matrix()` method.
   [(#3006)](https://github.com/PennyLaneAI/pennylane/pull/3006)
 
   ```pycon
@@ -291,6 +291,11 @@
   depth greater than 0. The `__repr__` for `Controlled` show `control_wires` instead of `wires`.
   [(#3013)](https://github.com/PennyLaneAI/pennylane/pull/3013)
 
+* Use `Operator.hash` instead of `Operator.matrix` to cache the eigendecomposition results in `Prod` and
+  `Sum` classes. When `Prod` and `Sum` operators have no overlapping wires, compute the eigenvalues
+  and the diagonalising gates using the factors/summands instead of using the full matrix.
+  [(#3022)](https://github.com/PennyLaneAI/pennylane/pull/3022)
+
 <h3>Breaking changes</h3>
 
 * Measuring an operator that might not be hermitian as an observable now raises a warning instead of an
@@ -323,7 +328,7 @@
 * Fixes a bug where the tape transform `single_qubit_fusion` computed wrong rotation angles
   for specific combinations of rotations.
   [(#3024)](https://github.com/PennyLaneAI/pennylane/pull/3024)
-    
+
 * Jax gradients now work with a QNode when the quantum function was transformed by `qml.simplify`.
   [(#3017)](https://github.com/PennyLaneAI/pennylane/pull/3017)
 

pennylane/ops/op_math/prod.py

@@ -167,6 +167,7 @@ def __init__(
         self.factors = factors
         self._wires = qml.wires.Wires.all_wires([f.wires for f in self.factors])
         self._hash = None
+        self._has_overlapping_wires = None
 
         if do_queue:
             self.queue()
@@ -226,6 +227,16 @@ def is_hermitian(self):
                 return False
         return all(op.is_hermitian for op in self.factors)
 
+    @property
+    def has_overlapping_wires(self) -> bool:
+        """Boolean expression that indicates if the factors have overlapping wires."""
+        if self._has_overlapping_wires is None:
+            wires = []
+            for op in self.factors:
+                wires.extend(list(op.wires))
+            self._has_overlapping_wires = len(wires) != len(set(wires))
+        return self._has_overlapping_wires
+
     def decomposition(self):
         r"""Decomposition of the product operator is given by each factor applied in succession.
 
@@ -250,14 +261,13 @@ def eigendecomposition(self):
             dict[str, array]: dictionary containing the eigenvalues and the
                 eigenvectors of the operator.
         """
-        Hmat = self.matrix()
-        Hmat = math.to_numpy(Hmat)
-        Hkey = tuple(Hmat.flatten().tolist())
-        if Hkey not in self._eigs:
+        if self.hash not in self._eigs:
+            Hmat = self.matrix()
+            Hmat = math.to_numpy(Hmat)
             w, U = np.linalg.eigh(Hmat)
-            self._eigs[Hkey] = {"eigvec": U, "eigval": w}
+            self._eigs[self.hash] = {"eigvec": U, "eigval": w}
 
-        return self._eigs[Hkey]
+        return self._eigs[self.hash]
 
     def diagonalizing_gates(self):
         r"""Sequence of gates that diagonalize the operator in the computational basis.
@@ -276,20 +286,31 @@ def diagonalizing_gates(self):
         Returns:
             list[.Operator] or None: a list of operators
         """
-
-        eigen_vectors = self.eigendecomposition["eigvec"]
-        return [qml.QubitUnitary(eigen_vectors.conj().T, wires=self.wires)]
+        if self.has_overlapping_wires:
+            eigen_vectors = self.eigendecomposition["eigvec"]
+            return [qml.QubitUnitary(eigen_vectors.conj().T, wires=self.wires)]
+        diag_gates = []
+        for factor in self.factors:
+            diag_gates.extend(factor.diagonalizing_gates())
+        return diag_gates
 
     def eigvals(self):
-        r"""Return the eigenvalues of the specified operator.
+        """Return the eigenvalues of the specified operator.
 
         This method uses pre-stored eigenvalues for standard observables where
         possible and stores the corresponding eigenvectors from the eigendecomposition.
 
         Returns:
             array: array containing the eigenvalues of the operator
         """
-        return self.eigendecomposition["eigval"]
+        if self.has_overlapping_wires:
+            return self.eigendecomposition["eigval"]
+        eigvals = [
+            qml.utils.expand_vector(factor.eigvals(), list(factor.wires), list(self.wires))
+            for factor in self.factors
+        ]
+
+        return qml.math.prod(eigvals, axis=0)
 
     def matrix(self, wire_order=None):
         """Representation of the operator as a matrix in the computational basis."""
@@ -423,7 +444,7 @@ def hash(self):
         return self._hash
 
 
-def _prod_sort(op_list, wire_map: dict = None):
+def _prod_sort(op_list, wire_map: dict = None) -> List[Operator]:
     """Insertion sort algorithm that sorts a list of product factors by their wire indices, taking
     into account the operator commutivity.
 

pennylane/ops/op_math/sum.py

@@ -153,6 +153,8 @@ def __init__(
         self._name = "Sum"
         self._id = id
         self.queue_idx = None
+        self._hash = None
+        self._has_overlapping_wires = None
 
         if len(summands) < 2:
             raise ValueError(f"Require at least two operators to sum; got {len(summands)}")
@@ -202,6 +204,16 @@ def is_hermitian(self):
         """If all of the terms in the sum are hermitian, then the Sum is hermitian."""
         return all(s.is_hermitian for s in self.summands)
 
+    @property
+    def has_overlapping_wires(self) -> bool:
+        """Boolean expression that indicates if the factors have overlapping wires."""
+        if self._has_overlapping_wires is None:
+            wires = []
+            for op in self.summands:
+                wires.extend(list(op.wires))
+            self._has_overlapping_wires = len(wires) != len(set(wires))
+        return self._has_overlapping_wires
+
     def terms(self):
         r"""Representation of the operator as a linear combination of other operators.
 
@@ -219,7 +231,7 @@ def terms(self):
 
     @property
     def eigendecomposition(self):
-        r"""Return the eigendecomposition of the matrix specified by the Hermitian observable.
+        r"""Return the eigendecomposition of the matrix specified by the operator.
 
         This method uses pre-stored eigenvalues for standard observables where
         possible and stores the corresponding eigenvectors from the eigendecomposition.
@@ -228,16 +240,15 @@ def eigendecomposition(self):
 
         Returns:
             dict[str, array]: dictionary containing the eigenvalues and the eigenvectors of the
-            operator
+                operator.
         """
-        Hmat = self.matrix()
-        Hmat = qml.math.to_numpy(Hmat)
-        Hkey = tuple(Hmat.flatten().tolist())
-        if Hkey not in self._eigs:
+        if self.hash not in self._eigs:
+            Hmat = self.matrix()
+            Hmat = math.to_numpy(Hmat)
             w, U = np.linalg.eigh(Hmat)
-            self._eigs[Hkey] = {"eigvec": U, "eigval": w}
+            self._eigs[self.hash] = {"eigvec": U, "eigval": w}
 
-        return self._eigs[Hkey]
+        return self._eigs[self.hash]
 
     def diagonalizing_gates(self):
         r"""Sequence of gates that diagonalize the operator in the computational basis.
@@ -256,9 +267,13 @@ def diagonalizing_gates(self):
         Returns:
             list[.Operator] or None: a list of operators
         """
-
-        eigen_vectors = self.eigendecomposition["eigvec"]
-        return [qml.QubitUnitary(eigen_vectors.conj().T, wires=self.wires)]
+        if self.has_overlapping_wires:
+            eigen_vectors = self.eigendecomposition["eigvec"]
+            return [qml.QubitUnitary(eigen_vectors.conj().T, wires=self.wires)]
+        diag_gates = []
+        for summand in self.summands:
+            diag_gates.extend(summand.diagonalizing_gates())
+        return diag_gates
 
     def eigvals(self):
         r"""Return the eigenvalues of the specified Hermitian observable.
@@ -269,7 +284,13 @@ def eigvals(self):
         Returns:
             array: array containing the eigenvalues of the Hermitian observable
         """
-        return self.eigendecomposition["eigval"]
+        if self.has_overlapping_wires:
+            return self.eigendecomposition["eigval"]
+        eigvals = [
+            qml.utils.expand_vector(summand.eigvals(), list(summand.wires), list(self.wires))
+            for summand in self.summands
+        ]
+        return qml.math.sum(eigvals, axis=0)
 
     def matrix(self, wire_order=None):
         r"""Representation of the operator as a matrix in the computational basis.
@@ -414,6 +435,14 @@ def simplify(self, cutoff=1.0e-12) -> "Sum":  # pylint: disable=arguments-differ
             else qml.Identity(self.wires[0]),
         )
 
+    @property
+    def hash(self):
+        if self._hash is None:
+            self._hash = hash(
+                (str(self.name), str([summand.hash for summand in _sum_sort(self.summands)]))
+            )
+        return self._hash
+
 
 class _SumSummandsGrouping:
     """Utils class used for grouping sum summands together."""
@@ -457,3 +486,35 @@ def get_summands(self, cutoff=1.0e-12):
                 new_summands.append(qml.s_prod(coeff, summand))
 
         return new_summands
+
+
+def _sum_sort(op_list, wire_map: dict = None) -> List[Operator]:
+    """Sort algorithm that sorts a list of sum summands by their wire indices.
+
+    Args:
+        op_list (List[.Operator]): list of operators to be sorted
+        wire_map (dict): Dictionary containing the wire values as keys and its indexes as values.
+            Defaults to None.
+
+    Returns:
+        List[.Operator]: sorted list of operators
+    """
+
+    if isinstance(op_list, tuple):
+        op_list = list(op_list)
+
+    def _sort_key(op) -> bool:
+        """Sorting key.
+
+        Args:
+            op (.Operator): Operator.
+
+        Returns:
+            int: Minimum wire value.
+        """
+        wires = op.wires
+        if wire_map is not None:
+            wires = wires.map(wire_map)
+        return np.min(wires)
+
+    return sorted(op_list, key=_sort_key)

tests/ops/op_math/test_prod.py

@@ -220,9 +220,7 @@ def test_eigen_caching(self):
         eig_vecs = eig_decomp["eigvec"]
         eig_vals = eig_decomp["eigval"]
 
-        eigs_cache = prod_op._eigs[
-            (1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, -1.0)
-        ]
+        eigs_cache = prod_op._eigs[prod_op.hash]
         cached_vecs = eigs_cache["eigvec"]
         cached_vals = eigs_cache["eigval"]
 
@@ -696,9 +694,7 @@ def test_eigen_caching(self):
 
         eig_vecs = eig_decomp["eigvec"]
         eig_vals = eig_decomp["eigval"]
-        eigs_cache = diag_prod_op._eigs[
-            (1.0, 0.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 1.0)
-        ]
+        eigs_cache = diag_prod_op._eigs[diag_prod_op.hash]
         cached_vecs = eigs_cache["eigvec"]
         cached_vals = eigs_cache["eigval"]
 
@@ -708,19 +704,7 @@ def test_eigen_caching(self):
     def test_diagonalizing_gates(self):
         """Test that the diagonalizing gates are correct."""
         diag_prod_op = Prod(qml.PauliZ(wires=0), qml.PauliZ(wires=1))
-        diagonalizing_gates = diag_prod_op.diagonalizing_gates()[0].matrix()
-        true_diagonalizing_gates = qnp.array(
-            (  # the gates that swap eigvals till they are ordered smallest --> largest
-                [
-                    [0.0, 1.0, 0.0, 0.0],
-                    [0.0, 0.0, 1.0, 0.0],
-                    [1.0, 0.0, 0.0, 0.0],
-                    [0.0, 0.0, 0.0, 1.0],
-                ]
-            )
-        )
-
-        assert np.allclose(diagonalizing_gates, true_diagonalizing_gates)
+        assert diag_prod_op.diagonalizing_gates() == []
 
 
 class TestSimplify:
@@ -1184,11 +1168,11 @@ def test_sorting_operators_with_wire_map(self):
             qml.PauliX(5),
             qml.Toffoli([2, "three", 4]),
             qml.CNOT([2, 5]),
-            qml.RX("test", 5),
+            qml.RX(1, 5),
             qml.PauliY(0),
-            qml.CRX("test", [0, 2]),
+            qml.CRX(1, [0, 2]),
             qml.PauliZ("three"),
-            qml.CRY("test", ["test", 2]),
+            qml.CRY(1, ["test", 2]),
         ]
         sorted_list = _prod_sort(op_list, wire_map={0: 0, "test": 1, 2: 2, "three": 3, 4: 4, 5: 5})
         final_list = [
@@ -1197,10 +1181,10 @@ def test_sorting_operators_with_wire_map(self):
             qml.Toffoli([2, "three", 4]),
             qml.PauliX(5),
             qml.CNOT([2, 5]),
-            qml.CRX("test", [0, 2]),
-            qml.CRY("test", ["test", 2]),
+            qml.CRX(1, [0, 2]),
+            qml.CRY(1, ["test", 2]),
             qml.PauliZ("three"),
-            qml.RX("test", 5),
+            qml.RX(1, 5),
         ]
 
         for op1, op2 in zip(final_list, sorted_list):