diff --git a/pennylane/math/quantum.py b/pennylane/math/quantum.py index b879f224871..ec27bed9d94 100644 --- a/pennylane/math/quantum.py +++ b/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 @@ def sqrt_matrix(density_matrix): 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.