Repository: PennyLaneAI/pennylane

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.


Discussion:

Resolving patch:
diff --git a/pennylane/ops/identity.py b/pennylane/ops/identity.py
--- a/pennylane/ops/identity.py
+++ b/pennylane/ops/identity.py
@@ -16,6 +16,7 @@
 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 @@ def compute_matrix():  # pylint: disable=arguments-differ
         """
         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
--- a/pennylane/ops/qubit/non_parametric_ops.py
+++ b/pennylane/ops/qubit/non_parametric_ops.py
@@ -21,6 +21,8 @@
 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 @@ def compute_matrix():  # pylint: disable=arguments-differ
         """
         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 @@ def compute_matrix():  # pylint: disable=arguments-differ
         """
         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 @@ def compute_matrix():  # pylint: disable=arguments-differ
         """
         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 @@ def compute_matrix():  # pylint: disable=arguments-differ
         """
         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
--- a/tests/ops/op_math/test_sprod.py
+++ b/tests/ops/op_math/test_sprod.py
@@ -299,18 +299,29 @@ def test_sprod_qubit_unitary(self):
         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
--- 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 @@
     (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 @@ def test_pow_independent_ops(self, op, n):
         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"),