Add a helper in math/quantum to calculate the matrix square root of sparse, hermitian matrices

pennylane/math/quantum.py

@@ -18,8 +18,11 @@
 import itertools
 from string import ascii_letters as ABC
 
+import scipy as sp
+import scipy.sparse.linalg as spla
 from autoray import numpy as np
 from numpy import float64  # pylint:disable=wrong-import-order
+from scipy.sparse import csc_matrix, issparse
 
 import pennylane as qml
 
@@ -972,6 +975,42 @@
     return vecs @ qml.math.diag(qml.math.sqrt(evs)) @ qml.math.conj(qml.math.transpose(vecs))
 
 
+def sqrt_matrix_sparse(density_matrix):
+    r"""Compute the square root matrix of a sparse density matrix where :math:`\rho = \sqrt{\rho} \times \sqrt{\rho}`
+
+    Args:
+        density_matrix (sparse): 2D sparse density matrix of the quantum system.
+
+    Returns:
+        (sparse): Square root of the density matrix. Even thought type as `csr_matrix` or `csc_matrix`, the output is not guaranteed to be sparse as well.
+    """
+    if not issparse(density_matrix):
+        raise TypeError(
+            "Only use this method for sparse matrices, or there will be an inevitable performance hit and divergence risk."
+        )
+    return _denman_beavers_iterations(density_matrix, max_iter=100, tol=1e-10)
+
+
+def _denman_beavers_iterations(mat, max_iter=100, tol=1e-10):
+    mat = csc_matrix(mat)
+    Ynew = mat
+
+    Znew = sp.sparse.eye(mat.shape[0]).tocsc()
+    for _ in range(max_iter):
+        Y = Ynew
+        Z = Znew
+        # TODO: implement the tol control logic
+        # basic idea: check norm_diff every 10 iter. If norm_diff < tol, break
+        Ynew = 0.5 * (Y + spla.inv(Z))
+        Znew = 0.5 * (Z + spla.inv(Y))
+        # TODO: common error catch to be here -- failure of spla.inv
+
+
+    X = spla.inv(Z)
+    B = Znew
+    return 2 * X - X @ B @ X
+
+
 def _compute_relative_entropy(rho, sigma, base=None):
     r"""
     Compute the quantum relative entropy of density matrix rho with respect to sigma.