QubitUnitary accepts sparse matrices as inputs

doc/releases/changelog-dev.md

@@ -10,6 +10,31 @@
 
 <h3>Improvements 🛠</h3>
 
+* `qml.QubitUnitary` now accepts sparse CSR matrices (from `scipy.sparse`). This allows efficient representation of large unitaries with mostly zero entries. Note that sparse unitaries are still in early development and may not support all features of their dense counterparts.
+  [(#6889)](https://github.com/PennyLaneAI/pennylane/pull/6889)
+
+  ```pycon
+  >>> import numpy as np
+  >>> import pennylane as qml
+  >>> import scipy as sp
+  >>> U_dense = np.eye(4)  # 2-wire identity
+  >>> U_sparse = sp.sparse.csr_matrix(U_dense)
+  >>> op = qml.QubitUnitary(U_sparse, wires=[0, 1])
+  >>> print(op.matrix())
+  <Compressed Sparse Row sparse matrix of dtype 'float64'
+          with 4 stored elements and shape (4, 4)>
+    Coords        Values
+    (0, 0)        1.0
+    (1, 1)        1.0
+    (2, 2)        1.0
+    (3, 3)        1.0
+  >>> op.matrix().toarray()
+  array([[1., 0., 0., 0.],
+        [0., 1., 0., 0.],
+        [0., 0., 1., 0.],
+        [0., 0., 0., 1.]])
+  ```
+
 * Add a decomposition for multi-controlled global phases into a one-less-controlled phase shift.
   [(#6936)](https://github.com/PennyLaneAI/pennylane/pull/6936)
 
@@ -86,7 +111,7 @@
   >>> print(new_circuit.diff_method)
   'parameter-shift'
   ```
-  
+
 * Devices can now configure whether or not ML framework data is sent to them
   via an `ExecutionConfig.convert_to_numpy` parameter. End-to-end jitting on
   `default.qubit` is used if the user specified a `jax.random.PRNGKey` as a seed.

pennylane/math/single_dispatch.py

@@ -19,8 +19,10 @@
 # pylint: disable=wrong-import-order
 import autoray as ar
 import numpy as np
+import scipy as sp
 from packaging.version import Version
 from scipy.linalg import block_diag as _scipy_block_diag
+from scipy.sparse.linalg import splu
 
 from .interface_utils import get_deep_interface
 from .utils import is_abstract
@@ -68,6 +70,68 @@ def _builtins_shape(x):
 ar.register_function("scipy", "ndim", np.ndim)
 
 
+# -------------------------------- SciPy Sparse --------------------------------- #
+# the following is required to ensure that general SciPy sparse matrices are
+# not automatically 'unwrapped' to dense NumPy arrays. Note that we assume
+# that whenever the backend is 'scipy', the input is a SciPy sparse matrix.
+
+
+def _det_sparse(x):
+    """Compute determinant of sparse matrices without densification"""
+
+    assert sp.sparse.issparse(x), TypeError(f"Expected SciPy sparse, got {type(x)}")
+
+    x = sp.sparse.csr_matrix(x)
+    if x.shape != (2, 2):
+        return _generic_sparse_det(x)
+
+    # Direct array access
+    indptr, indices, data = x.indptr, x.indices, x.data
+    values = {(i, j): 0.0 for i in range(2) for j in range(2)}
+    for i in range(2):
+        for j_idx in range(indptr[i], indptr[i + 1]):
+            j = indices[j_idx]
+            values[(i, j)] = data[j_idx]
+
+    return values[(0, 0)] * values[(1, 1)] - values[(0, 1)] * values[(1, 0)]
+
+
+def _generic_sparse_det(A):
+    """Compute the determinant of a sparse matrix using LU decomposition."""
+
+    assert hasattr(A, "tocsc"), TypeError(f"Expected SciPy sparse, got {type(A)}")
+
+    A_csc = A.tocsc()
+    lu = splu(A_csc)
+    U_diag = lu.U.diagonal()
+    det_A = np.prod(U_diag)
+    parity = _permutation_parity(lu.perm_r)
+    return det_A * parity
+
+
+def _permutation_parity(perm):
+    """Compute the parity of a permutation."""
+
+    parity = 1
+    visited = [False] * len(perm)
+    for i in range(len(perm)):
+        if not visited[i]:
+            cycle_length = 0
+            j = i
+            while not visited[j]:
+                visited[j] = True
+                j = perm[j]
+                cycle_length += 1
+
+            if cycle_length > 0:
+                parity *= (-1) ** (cycle_length - 1)
+    return parity
+
+
+ar.register_function("scipy", "linalg.det", _det_sparse)
+ar.register_function("scipy", "linalg.eigs", sp.sparse.linalg.eigs)
+ar.register_function("scipy", "trace", lambda x: x.trace())
+
 # -------------------------------- NumPy --------------------------------- #
 
 ar.register_function("numpy", "flatten", lambda x: x.flatten())

pennylane/math/utils.py

@@ -16,6 +16,7 @@
 # pylint: disable=wrong-import-order
 import autoray as ar
 import numpy as _np
+import scipy as sp
 
 # pylint: disable=import-outside-toplevel
 from autograd.numpy.numpy_boxes import ArrayBox
@@ -174,6 +175,9 @@ def convert_like(tensor1, tensor2):
         dev = tensor2.device
         return np.asarray(tensor1, device=dev, like=interface)
 
+    if interface == "scipy":
+        return sp.sparse.csr_matrix(tensor1)
+
     return np.asarray(tensor1, like=interface)
 
 

pennylane/ops/op_math/decompositions/single_qubit_unitary.py

@@ -14,8 +14,10 @@
 """Contains transforms and helpers functions for decomposing arbitrary unitary
 operations into elementary gates.
 """
+from functools import singledispatch
 
 import numpy as np
+import scipy as sp
 
 import pennylane as qml
 from pennylane import math
@@ -42,11 +44,16 @@ def _convert_to_su2(U, return_global_phase=False):
     with np.errstate(divide="ignore", invalid="ignore"):
         determinants = math.linalg.det(U)
     phase = math.angle(determinants) / 2
-    U = math.cast_like(U, determinants) * math.exp(-1j * math.cast_like(phase, 1j))[:, None, None]
+    U = (
+        math.cast_like(U, determinants) * math.exp(-1j * math.cast_like(phase, 1j))[:, None, None]
+        if not sp.sparse.issparse(U)
+        else U * math.exp(-1j * phase)
+    )
 
     return (U, phase) if return_global_phase else U
 
 
+@singledispatch
 def _zyz_get_rotation_angles(U):
     r"""Computes the rotation angles :math:`\phi`, :math:`\theta`, :math:`\omega`
     for a unitary :math:`U` that is :math:`SU(2)`
@@ -91,6 +98,51 @@ def _zyz_get_rotation_angles(U):
     return phis, thetas, omegas
 
 
+@_zyz_get_rotation_angles.register(sp.sparse.csr_matrix)
+def _zyz_get_rotation_angles_sparse(U):
+    r"""Computes the rotation angles :math:`\phi`, :math:`\theta`, :math:`\omega`
+    for a unitary :math:`U` that is :math:`SU(2)`, sparse case
+
+    Args:
+        U (array[complex]): A matrix that is :math:`SU(2)`
+
+    Returns:
+        tuple[array[float]]: A tuple containing the rotation angles
+            :math:`\phi`, :math:`\theta`, :math:`\omega`
+
+    """
+
+    assert sp.sparse.issparse(U), "Do not use this method if U is not sparse"
+
+    u00 = U[0, 0]
+    u01 = U[0, 1]
+    u10 = U[1, 0]
+
+    # For batched U or single U with non-zero off-diagonal, compute the
+    # normal decomposition instead
+    off_diagonal_elements = math.clip(math.abs(u01), 0, 1)
+    thetas = 2 * math.arcsin(off_diagonal_elements)
+
+    # Compute phi and omega from the angles of the top row; use atan2 to keep
+    # the angle within -np.pi and np.pi, and add very small value to the real
+    # part to avoid division by zero.
+    epsilon = 1e-64
+    angles_U00 = math.arctan2(math.imag(u00), math.real(u00) + epsilon)
+    angles_U10 = math.arctan2(math.imag(u10), math.real(u10) + epsilon)
+
+    phis = -angles_U10 - angles_U00
+    omegas = angles_U10 - angles_U00
+
+    phis, thetas, omegas = map(math.squeeze, [phis, thetas, omegas])
+
+    # Normalize the angles
+    phis = phis % (4 * np.pi)
+    thetas = thetas % (4 * np.pi)
+    omegas = omegas % (4 * np.pi)
+
+    return phis, thetas, omegas
+
+
 def _rot_decomposition(U, wire, return_global_phase=False):
     r"""Compute the decomposition of a single-qubit matrix :math:`U` in terms of
     elementary operations, as a single :class:`.RZ` gate or a :class:`.Rot` gate.
@@ -167,7 +219,7 @@ def _get_single_qubit_rot_angles_via_matrix(
     of the matrix of the target operation using ZYZ rotations.
     """
     # Cast to batched format for more consistent code
-    U = math.expand_dims(U, axis=0) if len(U.shape) == 2 else U
+    U = math.expand_dims(U, axis=0) if len(U.shape) == 2 and not sp.sparse.issparse(U) else U
 
     # Convert to SU(2) format and extract global phase
     U_su2, global_phase = _convert_to_su2(U, return_global_phase=True)

pennylane/ops/op_math/decompositions/two_qubit_unitary.py

@@ -17,6 +17,7 @@
 import warnings
 
 import numpy as np
+import scipy as sp
 
 import pennylane as qml
 from pennylane import math
@@ -102,6 +103,9 @@
 )
 
 
+global_arrays_name = ["E", "Edag", "CNOT01", "CNOT10", "SWAP", "S_SX", "v_one_cnot", "q_one_cnot"]
+
+
 def _convert_to_su4(U):
     r"""Convert a 4x4 matrix to :math:`SU(4)`.
 
@@ -208,7 +212,7 @@
             [0],  # arbitrary value for x
         )
 
     return math.convert_like(A, U), math.convert_like(B, U)
 
 
 def _extract_su2su2_prefactors(U, V):
@@ -621,6 +625,11 @@
     _check_differentiability_warning(U)
     # First, we note that this method works only for SU(4) gates, meaning that
     # we need to rescale the matrix by its determinant.
+    if sp.sparse.issparse(U):
+        # Convert all the global elements to sparse matrices in-place
+        for name in global_arrays_name:
+            array = globals()[name]
+            globals()[name] = sp.sparse.csr_matrix(array)
     U = _convert_to_su4(U)
 
     # The next thing we will do is compute the number of CNOTs needed, as this affects
@@ -631,7 +640,7 @@
     with qml.QueuingManager.stop_recording():
         if qml.math.is_abstract(U):
             decomp = _decomposition_3_cnots(U, wires)
         elif num_cnots == 0:
             decomp = _decomposition_0_cnots(U, wires)
         elif num_cnots == 1:
             decomp = _decomposition_1_cnot(U, wires)

pennylane/ops/qubit/matrix_ops.py

@@ -21,7 +21,9 @@
 from typing import Optional, Union
 
 import numpy as np
+import scipy as sp
 from scipy.linalg import fractional_matrix_power
+from scipy.sparse import csr_matrix
 
 import pennylane as qml
 from pennylane import numpy as pnp
@@ -78,6 +80,10 @@ class QubitUnitary(Operation):
     r"""QubitUnitary(U, wires)
     Apply an arbitrary unitary matrix with a dimension that is a power of two.
 
+    .. warning::
+
+        The sparse matrix representation of QubitUnitary is still under development. Currently we only support a limited set of interfaces that preserve the sparsity of the matrix, including ..method::`adjoint`, ..method::`pow`, ..method::`compute_sparse_matrix` and ..method::`compute_decomposition`. Differentiability is not supported for sparse matrices.
+
     **Details:**
 
     * Number of wires: Any (the operation can act on any number of wires)
@@ -86,7 +92,7 @@ class QubitUnitary(Operation):
     * Gradient recipe: None
 
     Args:
-        U (array[complex]): square unitary matrix
+        U (array[complex] or csr_matrix): square unitary matrix
         wires (Sequence[int] or int): the wire(s) the operation acts on
         id (str): custom label given to an operator instance,
             can be useful for some applications where the instance has to be identified
@@ -121,7 +127,7 @@ class QubitUnitary(Operation):
 
     def __init__(
         self,
-        U: TensorLike,
+        U: Union[TensorLike, csr_matrix],
         wires: WiresLike,
         id: Optional[str] = None,
         unitary_check: bool = False,
@@ -135,19 +141,16 @@ def __init__(
         if len(U_shape) not in {2, 3} or U_shape[-2:] != (dim, dim):
             raise ValueError(
                 f"Input unitary must be of shape {(dim, dim)} or (batch_size, {dim}, {dim}) "
-                f"to act on {len(wires)} wires."
+                f"to act on {len(wires)} wires. Got shape {U_shape} instead."
             )
 
+        # If the matrix is sparse, we need to convert it to a csr_matrix
+        if sp.sparse.issparse(U):
+            U = U.tocsr()
+
         # Check for unitarity; due to variable precision across the different ML frameworks,
         # here we issue a warning to check the operation, instead of raising an error outright.
-        if unitary_check and not (
-            qml.math.is_abstract(U)
-            or qml.math.allclose(
-                qml.math.einsum("...ij,...kj->...ik", U, qml.math.conj(U)),
-                qml.math.eye(dim),
-                atol=1e-6,
-            )
-        ):
+        if unitary_check and not self._unitary_check(U, dim):
             warnings.warn(
                 f"Operator {U}\n may not be unitary. "
                 "Verify unitarity of operation, or use a datatype with increased precision.",
@@ -156,6 +159,18 @@ def __init__(
 
         super().__init__(U, wires=wires, id=id)
 
+    @staticmethod
+    def _unitary_check(U, dim):
+        if isinstance(U, csr_matrix):
+            U_dagger = U.conjugate().transpose()
+            identity = sp.sparse.eye(dim, format="csr")
+            return sp.sparse.linalg.norm(U @ U_dagger - identity) < 1e-10
+        return qml.math.allclose(
+            qml.math.einsum("...ij,...kj->...ik", U, qml.math.conj(U)),
+            qml.math.eye(dim),
+            atol=1e-6,
+        )
+
     @staticmethod
     def compute_matrix(U: TensorLike):  # pylint: disable=arguments-differ
         r"""Representation of the operator as a canonical matrix in the computational basis (static method).
@@ -180,6 +195,26 @@ def compute_matrix(U: TensorLike):  # pylint: disable=arguments-differ
         """
         return U
 
+    @staticmethod
+    def compute_sparse_matrix(U: TensorLike):  # pylint: disable=arguments-differ
+        r"""Representation of the operator as a sparse matrix.
+
+        Args:
+            U (tensor_like): unitary matrix
+
+        Returns:
+            csr_matrix: sparse matrix representation
+
+        **Example**
+
+        >>> U = np.array([[0.98877108+0.j, 0.-0.14943813j], [0.-0.14943813j, 0.98877108+0.j]])
+        >>> qml.QubitUnitary.compute_sparse_matrix(U)
+        <2x2 sparse matrix of type '<class 'numpy.complex128'>'
+            with 2 stored elements in Compressed Sparse Row format>
+        """
+        U = qml.math.asarray(U, like="numpy")
+        return sp.sparse.csr_matrix(U)
+
     @staticmethod
     def compute_decomposition(U: TensorLike, wires: WiresLike):
         r"""Representation of the operator as a product of other operators (static method).
@@ -236,10 +271,17 @@ def has_decomposition(self) -> bool:
 
     def adjoint(self) -> "QubitUnitary":
         U = self.matrix()
+        if isinstance(U, csr_matrix):
+            adjoint_sp_mat = U.conjugate().transpose()
+            # Note: it is necessary to explicitly cast back to csr, or it will be come csc
+            return QubitUnitary(csr_matrix(adjoint_sp_mat), wires=self.wires)
         return QubitUnitary(qml.math.moveaxis(qml.math.conj(U), -2, -1), wires=self.wires)
 
     def pow(self, z: Union[int, float]):
         mat = self.matrix()
+        if isinstance(mat, csr_matrix):
+            pow_mat = sp.sparse.linalg.matrix_power(mat, z)
+            return [QubitUnitary(pow_mat, wires=self.wires)]
         if isinstance(z, int) and qml.math.get_deep_interface(mat) != "tensorflow":
             pow_mat = qml.math.linalg.matrix_power(mat, z)
         elif self.batch_size is not None or qml.math.shape(z) != ():