Repository: PennyLaneAI/pennylane (python)

Issue:
Add sparse matrix representations for Pauli Operators
While the Pauli operators (`PauliX`, `PauliY`, `PauliZ`, `Identity`) themselves are fairly dense, they are often multiplied together to form large tensors. To assist with the formation of sparse matrices for large combination operations, we should add sparse matrices for the basic Pauli operators.

Following Issue #2881, these `compute_sparse_matrix` methods should accept a `format` keyword.

Pull request description:
Support Sparse Matrix representation for Observables and Arithmetic Ops (1/3)
**Context:**
We want to support sparse matrix representation for Symbolic operators, these depend on the sparse matrix attributes of their factors/summands/base operators. We must then also support sparse matrix representations for a few other operators; we choose to support observables as these have useful applications (eg. Quantum Chemistry). 

**Description of the Change:**
Implement the `sparse_matrix()` property for some well known PL observables. Add a sparse expand method to handle varying wire labels. 

**Benefits:**
Performance boost when computing expectation values / other measurements using operator arithmetic or observables. 

Closes #2882 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index cbe4b45ba..f16ecddac 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -112,6 +112,9 @@
   >>> qml.simplify(adj_op)
   RX(-1, wires=[0])
   ```
+  
+* Added `sparse_matrix()` support for single qubit observables
+  [(#2964)](https://github.com/PennyLaneAI/pennylane/pull/2964)
 
 * Added the `qml.is_hermitian` and `qml.is_unitary` function checks.
   [(#2960)](https://github.com/PennyLaneAI/pennylane/pull/2960)
@@ -152,6 +155,7 @@ Albert Mitjans Coma,
 Rashid N H M,
 Zeyue Niu,
 Mudit Pandey,
-Antal Száva,
+Jay Soni,
+Antal Száva
 Cody Wang,
 David Wierichs
diff --git a/pennylane/ops/identity.py b/pennylane/ops/identity.py
index 228092153..5b7dac068 100644
--- a/pennylane/ops/identity.py
+++ b/pennylane/ops/identity.py
@@ -16,6 +16,7 @@ This module contains the Identity operation that is common to both
 cv and qubit computing paradigms in PennyLane.
 """
 import numpy as np
+from scipy import sparse
 
 from pennylane.operation import CVObservable, Operation
 
@@ -91,6 +92,10 @@ class Identity(CVObservable, Operation):
         """
         return np.eye(2)
 
+    @staticmethod
+    def compute_sparse_matrix(*params, **hyperparams):
+        return sparse.csr_matrix([[1, 0], [0, 1]])
+
     @staticmethod
     def _heisenberg_rep(p):
         return np.array([1, 0, 0])
diff --git a/pennylane/ops/qubit/non_parametric_ops.py b/pennylane/ops/qubit/non_parametric_ops.py
index ab01efd9d..c595ba7ce 100644
--- a/pennylane/ops/qubit/non_parametric_ops.py
+++ b/pennylane/ops/qubit/non_parametric_ops.py
@@ -21,6 +21,8 @@ import warnings
 from copy import copy
 
 import numpy as np
+
+from scipy import sparse
 from scipy.linalg import block_diag
 
 import pennylane as qml
@@ -76,6 +78,11 @@ class Hadamard(Observable, Operation):
         """
         return np.array([[INV_SQRT2, INV_SQRT2], [INV_SQRT2, -INV_SQRT2]])
 
+    @staticmethod
+    def compute_sparse_matrix(*params, **hyperparams):
+        """Compute the sparse matrix representation"""
+        return sparse.csr_matrix([[INV_SQRT2, INV_SQRT2], [INV_SQRT2, -INV_SQRT2]])
+
     @staticmethod
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
@@ -214,6 +221,10 @@ class PauliX(Observable, Operation):
         """
         return np.array([[0, 1], [1, 0]])
 
+    @staticmethod
+    def compute_sparse_matrix(*params, **hyperparams):
+        return sparse.csr_matrix([[0, 1], [1, 0]])
+
     @staticmethod
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
@@ -358,6 +369,10 @@ class PauliY(Observable, Operation):
         """
         return np.array([[0, -1j], [1j, 0]])
 
+    @staticmethod
+    def compute_sparse_matrix(*params, **hyperparams):
+        return sparse.csr_matrix([[0, -1j], [1j, 0]])
+
     @staticmethod
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
@@ -500,6 +515,10 @@ class PauliZ(Observable, Operation):
         """
         return np.array([[1, 0], [0, -1]])
 
+    @staticmethod
+    def compute_sparse_matrix(*params, **hyperparams):
+        return sparse.csr_matrix([[1, 0], [0, -1]])
+
     @staticmethod
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).


Test patch:
diff --git a/tests/ops/op_math/test_sprod.py b/tests/ops/op_math/test_sprod.py
index 668d5c181..cdbb0c113 100644
--- a/tests/ops/op_math/test_sprod.py
+++ b/tests/ops/op_math/test_sprod.py
@@ -299,18 +299,29 @@ class TestMatrix:
         true_mat = 42 * U
         assert np.allclose(mat, true_mat)
 
-    # TODO[Jay]: remove xfail once there is support for sparse matrices for most operations
-    @pytest.mark.xfail
-    @pytest.mark.parametrize("scalar, op", ops)
-    def test_sparse_matrix(self, op, scalar):
+    sparse_ops = (
+        qml.Identity(wires=0),
+        qml.PauliX(wires=0),
+        qml.PauliY(wires=0),
+        qml.PauliZ(wires=0),
+        qml.Hadamard(wires=0),
+    )
+
+    @pytest.mark.parametrize("op", sparse_ops)
+    def test_sparse_matrix(self, op):
         """Test the sparse_matrix representation of scaled ops."""
+        scalar = 1 + 2j
         sprod_op = SProd(scalar, op)
         sparse_matrix = sprod_op.sparse_matrix()
+        sparse_matrix.sort_indices()
 
         expected_sparse_matrix = scalar * op.matrix()
         expected_sparse_matrix = csr_matrix(expected_sparse_matrix)
+        expected_sparse_matrix.sort_indices()
 
-        assert np.allclose(sparse_matrix.todense(), expected_sparse_matrix.todense())
+        assert type(sparse_matrix) == type(expected_sparse_matrix)
+        assert all(sparse_matrix.data == expected_sparse_matrix.data)
+        assert all(sparse_matrix.indices == expected_sparse_matrix.indices)
 
     def test_sparse_matrix_sparse_hamiltonian(self):
         """Test the sparse_matrix representation of scaled ops."""
diff --git a/tests/ops/qubit/test_non_parametric_ops.py b/tests/ops/qubit/test_non_parametric_ops.py
index 21cfe8d7c..444d345b6 100644
--- a/tests/ops/qubit/test_non_parametric_ops.py
+++ b/tests/ops/qubit/test_non_parametric_ops.py
@@ -15,13 +15,14 @@
 Unit tests for the available non-parametric qubit operations
 """
 import itertools
-import re
 import pytest
 import copy
 import numpy as np
+from scipy.sparse import csr_matrix
 from scipy.stats import unitary_group
 
 import pennylane as qml
+from pennylane import math
 from pennylane.wires import Wires
 
 from gate_data import I, X, Y, Z, H, CNOT, SWAP, ISWAP, SISWAP, CZ, S, T, CSWAP, Toffoli, ECR
@@ -42,6 +43,14 @@ NON_PARAMETRIZED_OPERATIONS = [
     (qml.ECR, ECR),
 ]
 
+SPARSE_MATRIX_SUPPORTED_OPERATIONS = (
+    (qml.Identity(wires=0), I),
+    (qml.Hadamard(wires=0), H),
+    (qml.PauliZ(wires=0), Z),
+    (qml.PauliX(wires=0), X),
+    (qml.PauliY(wires=0), Y),
+)
+
 
 class TestOperations:
     @pytest.mark.parametrize("op_cls, mat", NON_PARAMETRIZED_OPERATIONS)
@@ -1160,6 +1169,17 @@ class TestPowMethod:
         assert op.pow(n)[0].__class__ is op.__class__
 
 
+class TestSparseMatrix:
+    @pytest.mark.parametrize("op, mat", SPARSE_MATRIX_SUPPORTED_OPERATIONS)
+    def test_sparse_matrix(self, op, mat):
+        expected_sparse_mat = csr_matrix(mat)
+        sparse_mat = op.sparse_matrix()
+
+        assert type(sparse_mat) == type(expected_sparse_mat)
+        assert all(sparse_mat.data == expected_sparse_mat.data)
+        assert all(sparse_mat.indices == expected_sparse_mat.indices)
+
+
 label_data = [
     (qml.Identity(0), "I", "I"),
     (qml.Hadamard(0), "H", "H"),