Repository: PennyLaneAI/pennylane
Pull request #1596 (2021-08-27 06:51:19)

Vectorize and remove for loop in sparse expvals
**Context:**

In #1551, support for computing the expectation values of Hamiltonians using sparse methods was directly added to `default.qubit`. However, two improvements were noticed:

- The `for` loop over the Hamiltonian observables and coefficients could be replaced with vectorization

- When not in backprop mode, the full sparse matrix could be computed directly.

**Description of the Change:**

- The `for` loop discussed above has been removed and replaced with vectorization.

- When non-backprop mode is detected, we simply re-use the same logic that is currently used for `qml.SparseHamiltonian`.

**Benefits:**

- Vectorizing the for loop allows us to remove the casting of Hamiltonian coefficients to NumPy arrays, that we were performing as a bugfix.

- `qml.SparseHamiltonian` and `qml.Hamiltonian` now share the same code for expvals when using parameter-shift.

- Faster execution as opposed to backprop in master

  ```
  Backprop:
    Using ExpvalCost  : 0.9697349071502686     (1x)
    expval(H), master : 0.2827181816101074   (3.4x)
    expval(H), this PR: 0.08908772468566895 (10.9x)
  ```

- Faster execution as opposed to parameter-shift in master

  ```
  Parameter-shift:
    Using ExpvalCost  : 6.549182176589966    (1x)
    expval(H), master : 2.6744325160980225 (2.4x)
    expval(H), this PR: 1.9294066429138184 (3.4x)
  ```

**Possible Drawbacks:**

- While improvements to the parameter-shift pipeline in `default.qubit` are nice, we plan to transition to `lightning.qubit` by default, and likely restrict `default.qubit` to purely backprop going forward.

- Is the `SparseHamiltonian` class deprecated? Should we remove it?

**Related GitHub Issues:** n/a



Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 1fb629a1d..d653dc113 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -258,6 +258,7 @@
 * Hamiltonians are now natively supported on the `default.qubit` device if `shots=None`. 
   This makes VQE workflows a lot faster in some cases.
   [(#1551)](https://github.com/PennyLaneAI/pennylane/pull/1551)
+  [(#1596)](https://github.com/PennyLaneAI/pennylane/pull/1596)
 
 * A gradient recipe for Hamiltonian coefficients has been added. This makes it possible 
   to compute parameter-shift gradients of these coefficients on devices that natively 
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index b335bdaf8..cb53a6ca6 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -467,43 +467,52 @@ class DefaultQubit(QubitDevice):
         Returns:
             float: returns the expectation value of the observable
         """
-        if observable.name == "SparseHamiltonian":
-            if self.shots is not None:
-                raise DeviceError("SparseHamiltonian must be used with shots=None")
-
-            ev = coo_matrix.dot(
-                coo_matrix(self._conj(self.state)),
-                coo_matrix.dot(
-                    observable.matrix, coo_matrix(self.state.reshape(len(self.state), 1))
-                ),
-            )
-
-            return np.real(ev.toarray()[0])
-
         # intercept Hamiltonians here; in future, we want a logic that handles
         # general observables that do not define eigenvalues
-        if observable.name == "Hamiltonian":
-            assert self.shots is None, "Hamiltonian must be used with shots=None"
-
-            # compute  <psi| H |psi> via sum_i coeff_i * <psi| PauliWord |psi> using a sparse
-            # representation of the Pauliword
-            res = qml.math.cast(qml.math.convert_like(0.0, observable.coeffs), dtype=complex)
-            # note: it is important that we use the Hamiltonian's data and not the coeffs attribute
-            for op, coeff in zip(observable.ops, observable.data):
-                # extract a scipy.sparse.coo_matrix representation of this Pauli word
-                coo = qml.operation.Tensor(op).sparse_matrix(wires=self.wires)
-                for idx_row, idx_col, entry in zip(coo.row, coo.col, coo.data):
-                    # while "entry" is not differentiable, it will be parsed during multiplication
-                    product = self._conj(self.state)[idx_row] * entry * self.state[idx_col]
-
-                    # todo: remove this hack that avoids errors when attempting to multiply
-                    # a nontrainable qml.tensor to a trainable Arraybox
-                    if isinstance(coeff, qml.numpy.tensor) and not coeff.requires_grad:
-                        coeff = qml.math.toarray(coeff)
-
-                    res = qml.math.convert_like(res, product) + (
-                        qml.math.cast(qml.math.convert_like(coeff, product), "complex128") * product
+        if observable.name in ("Hamiltonian", "SparseHamiltonian"):
+            assert self.shots is None, f"{observable.name} must be used with shots=None"
+
+            backprop_mode = not isinstance(self.state, np.ndarray)
+
+            if backprop_mode:
+                # We must compute the expectation value assuming that the Hamiltonian
+                # coefficients *and* the quantum states are tensor objects.
+
+                # Compute  <psi| H |psi> via sum_i coeff_i * <psi| PauliWord |psi> using a sparse
+                # representation of the Pauliword
+                res = qml.math.cast(qml.math.convert_like(0.0, observable.data), dtype=complex)
+
+                # Note: it is important that we use the Hamiltonian's data and not the coeffs attribute.
+                # This is because the .data attribute may be 'unwrapped' as required by the interfaces,
+                # whereas the .coeff attribute will always be the same input dtype that the user provided.
+                for op, coeff in zip(observable.ops, observable.data):
+
+                    # extract a scipy.sparse.coo_matrix representation of this Pauli word
+                    coo = qml.operation.Tensor(op).sparse_matrix(wires=self.wires)
+                    Hmat = qml.math.cast(qml.math.convert_like(coo.data, self.state), "complex128")
+
+                    product = (
+                        qml.math.gather(qml.math.conj(self.state), coo.row)
+                        * Hmat
+                        * qml.math.gather(self.state, coo.col)
                     )
+                    c = qml.math.cast(qml.math.convert_like(coeff, product), "complex128")
+                    res = qml.math.convert_like(res, product) + qml.math.sum(c * product)
+
+            else:
+                # Coefficients and the state are not trainable, we can be more
+                # efficient in how we compute the Hamiltonian sparse matrix.
+
+                if observable.name == "Hamiltonian":
+                    Hmat = qml.utils.sparse_hamiltonian(observable, wires=self.wires)
+                elif observable.name == "SparseHamiltonian":
+                    Hmat = observable.matrix
+
+                res = coo_matrix.dot(
+                    coo_matrix(qml.math.conj(self.state)),
+                    coo_matrix.dot(Hmat, coo_matrix(self.state.reshape(len(self.state), 1))),
+                ).toarray()[0]
+
             return qml.math.real(res)
 
         return super().expval(observable, shot_range=shot_range, bin_size=bin_size)
diff --git a/pennylane/utils.py b/pennylane/utils.py
index 3e2134db4..8e69feb90 100644
--- a/pennylane/utils.py
+++ b/pennylane/utils.py
@@ -150,7 +150,7 @@ def sparse_hamiltonian(H, wires=None):
     n = len(wires)
     matrix = scipy.sparse.coo_matrix((2 ** n, 2 ** n), dtype="complex128")
 
-    coeffs = qml.math.toarray(H.coeffs)
+    coeffs = qml.math.toarray(H.data)
 
     for coeff, op in zip(coeffs, H.ops):
 


Test patch:
diff --git a/tests/ops/test_sparse.py b/tests/ops/test_sparse.py
index 144f57888..05a0b1d36 100644
--- a/tests/ops/test_sparse.py
+++ b/tests/ops/test_sparse.py
@@ -173,5 +173,5 @@ class TestSparse:
 
         dev = qml.device("default.qubit", wires=1, shots=1)
 
-        with pytest.raises(DeviceError, match="SparseHamiltonian must be used with shots=None"):
+        with pytest.raises(AssertionError, match="SparseHamiltonian must be used with shots=None"):
             dev.expval(qml.SparseHamiltonian(hamiltonian, [0]))[0]