Repository: PennyLaneAI/pennylane
Pull request #6976 (2025-02-18 21:40:35)

Add a helper in `math/quantum` to calculate the matrix square root of sparse, hermitian matrices
**Context:**
This need arose from https://github.com/PennyLaneAI/pennylane/pull/6963 where we realized that `scipy.sparse` could NOT help with a sparsity-preserved `sqrtm` as a sparse counterpart of its own `scipy.linalg.sqrtm` using Schur decomposition. Soon we realized that there could be a long list of lost functionality, e.g. `eigh`, `sqrtm` etc. We would like to update step by step such that when we implement the sparse representation for various matrices we do not have to convert them into dense repr, but this might involve quite some effort considering performance benchmark and lack of references. Therefore, we need some temporary solution that helps with specific `Operator` such as `BlockEncode` to avoid an explicit cast back to dense, thought the performance might vary.

For the `BlockEncode` considered in this PR, it specifically used `sqrtm` for a special class of matrices $Id - A^\dagger A$ where $A$ is the input matrix to encode into quantum circuits. This is very close to our own `sqrt_matrix` in `qml.math.quantum` that deals with the square root computation of a density matrix; our previous dense-matrix algorithm also made this assumption and directly called `sqrt_matrix` in the needed place. Hence we put forward a commonly used iterative method here suitable for similar Hermitian matrices to have square roots.

A quick introduction of what is the Denman–Beavers iteration used by this PR can be found at [Wikipedia](https://en.wikipedia.org/wiki/Square_root_of_a_matrix#By_Denman%E2%80%93Beavers_iteration).

**Description of the Change:**
Provide a new helper in `quantum.py`

**Benefits:**
Enable the sparse version of `BlockEncode`

**Possible Drawbacks:**
This method might have poor performance if used improperly. Empirically, do not go over 100by100 or it will cost at least several seconds to converge.
It might not converge when the input matrix is away from positive-definite Hermitian, unlike our existing `sqrt_matrix` that works for dense matrices, which in general can handle any Hermitian matrices)

**Related GitHub Issues:**
[sc-84693]

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6976?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `86.11111%` with `5 lines` in your changes missing coverage. Please review.
> Project coverage is 99.58%. Comparing base [(`2f5a6ad`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/2f5a6ad4fc2fbb22e750cd8b4a761d4a589af834?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`48ea790`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/48ea790320bfdb1dfbe81dc317eed8f6fd3bb90b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 1 commits behind head on master.

| [Files with missing lines](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6976?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/math/quantum.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6976?src=pr&el=tree&filepath=pennylane%2Fmath%2Fquantum.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvcXVhbnR1bS5weQ==) | 86.11% | [5 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6976?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #6976      +/-   ##
==========================================
- Coverage   99.59%   99.58%   -0.02%     
==========================================
  Files         483      483              
  Lines       46119    46155      +36     
==========================================
+ Hits        45934    45965      +31     
- Misses        185      190       +5     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6976?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

JerryChen97: Some naive benmarks with few tuned parameters over mat size and nnz (density)
![DenmanBeavers_time_00](https://github.com/user-attachments/assets/13774876-b281-4056-abde-e0f40ef4080d)
![DenmanBeavers_time_Heatmap_00](https://github.com/user-attachments/assets/47709e23-728a-4864-9b29-6bd938c9906c)

JerryChen97: Note how expensive it could be when the matrix sizes go over hundred. We might want to also warn users on any usage beyond this empirical threshold.
dwierichs: >Note how expensive it could be when the matrix sizes go over hundred. We might want to also warn users on any usage beyond this empirical threshold.

I'm not sure we need to do this. I think a performance warning is useful if it gives actionable advice like "change data types because what you're doing is unnecessarily slow" or "this method was not built for your inputs". So if the warning is to actually direct a user back to dense matrices, we could do that (if those are actually any faster :sweat_smile: ), but otherwise I'd consider the warning "your problem is large so it takes some time to compute" not to be super valuable.
I might be unaware of some convention we have for this, though, and @mlxd will know better :)
JerryChen97: @mlxd 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 086d1e9ad..12070d8ad 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,9 @@
 
 <h3>New features since last release</h3>
 
+* Added method `qml.math.sqrt_matrix_sparse` to compute the square root of a sparse Hermitian matrix.
+  [(#6976)](https://github.com/PennyLaneAI/pennylane/pull/6976)
+
 * Added a class `qml.capture.transforms.MergeRotationsInterpreter` that merges rotation operators
   following the same API as `qml.transforms.optimization.merge_rotations` when experimental program capture is enabled.
   [(#6957)](https://github.com/PennyLaneAI/pennylane/pull/6957)
diff --git a/pennylane/math/__init__.py b/pennylane/math/__init__.py
index 930560eab..048aef21a 100644
--- a/pennylane/math/__init__.py
+++ b/pennylane/math/__init__.py
@@ -76,6 +76,7 @@ from .quantum import (
     reduce_statevector,
     relative_entropy,
     sqrt_matrix,
+    sqrt_matrix_sparse,
     vn_entropy,
     vn_entanglement_entropy,
     max_entropy,
diff --git a/pennylane/math/quantum.py b/pennylane/math/quantum.py
index b879f2248..3750da1a6 100644
--- a/pennylane/math/quantum.py
+++ b/pennylane/math/quantum.py
@@ -18,8 +18,11 @@ import functools
 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,118 @@ 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(sparse_matrix):
+    r"""Compute the square root matrix of a positive-definite Hermitian matrix where :math:`\rho = \sqrt{\rho} \times \sqrt{\rho}`
+
+    Args:
+        sparse_matrix (sparse): 2D sparse matrix of the quantum system.
+
+    Returns:
+       (sparse): Square root of the sparse matrix. Even for data types like `csr_matrix` or `csc_matrix`, the output matrix is not guaranteed to be sparse as well.
+
+    """
+    if not issparse(sparse_matrix):
+        raise TypeError(
+            f"sqrt_matrix_sparse currently only supports scipy.sparse matrices, but received {type(sparse_matrix)}. "
+        )
+    if sparse_matrix.nnz == 0:
+        return sparse_matrix
+    # NOTE: the following steps should be re-visited in the future to establish
+    # better understanding and control over the heuristics we chose
+    # 1. choice of max iteration and tolerance for denman beavers, sc-85713
+    # 2. different methods for sparse matrix square root, sc-85710
+    return _denman_beavers_iterations(sparse_matrix, max_iter=100, tol=1e-10)
+
+
+def _inv_newton(M, guess):
+    """
+    Compute the inverse of a matrix using Newton's method.
+
+    Args:
+        M (array-like): The matrix to be inverted.
+        guess (array-like): An initial guess for the inverse of the matrix.
+
+    Returns:
+        array-like: An improved estimate of the inverse of the matrix.
+    """
+    return 2 * guess - guess @ M @ guess
+
+
+def _denman_beavers_iterations(mat, max_iter=100, tol=1e-13):
+    """Compute matrix square root using the Denman-Beavers iteration.
+
+    The Denman–Beavers iteration was introduced by E. D. Denman and A. N. Beavers in 1976
+    and stems from Newton-type methods originally used to compute the matrix sign function.
+    In this adaptation for matrix square roots, two matrices (Y and Z) are refined in each
+    step until convergence. This technique is often effective for sparse or structured
+    matrices, particularly those that are positive semidefinite or invertible.
+
+    Args:
+        mat (sparse): Sparse input matrix
+        max_iter (int): Maximum number of iterations
+        tol (float): Convergence tolerance (absolute tolerance). Measured using the Frobenius norm of the difference between input mat and the square of the output.
+
+    Returns:
+        scipy.sparse.spmatrix: Square root of the input matrix
+
+    Raises:
+        LinAlgError: If matrix inversion fails
+        ValueError: If NaN values or overflow are encountered during computation
+    """
+    try:
+        mat = csc_matrix(mat)
+        Y = mat
+        Z = sp.sparse.eye(mat.shape[0], format="csc")
+
+        # Keep track of previous iteration for convergence check
+        Y_prev = None
+        norm_diff = None
+
+        for iter_num in range(max_iter):
+            # Compute next iteration
+            if iter_num < 2:
+                Zinv = spla.inv(Z) if iter_num > 0 else Z
+                Yinv = spla.inv(Y)
+            else:
+                # Take Newton step
+                Zinv = _inv_newton(Z, Zinv)
+                Yinv = _inv_newton(Y, Yinv)
+
+            Y = 0.5 * (Y + Zinv)
+            Z = 0.5 * (Z + Yinv)
+
+            # Check for NaN or infinite values
+            if not (np.all(np.isfinite(Y.data)) and np.all(np.isfinite(Z.data))):
+                raise ValueError(
+                    "Invalid values encountered during computation: nan or inf"
+                    f"Input matrix: {mat.toarray()}"
+                )
+
+            # Check convergence every 10 iterations
+            if iter_num % 10 == 0 and iter_num > 0:
+                if Y_prev is not None:
+                    # Compute Frobenius norm of difference
+                    diff = Y - Y_prev
+                    norm_diff = spla.norm(diff)
+                    if norm_diff < tol:
+                        break
+                Y_prev = Y.copy()
+
+        numerical_error = spla.norm((Y @ Y - mat))
+        if (norm_diff and norm_diff > tol) or numerical_error > tol:
+            raise ValueError(
+                f"Convergence threshold not reached after {max_iter} iterations, "
+                f"with final norm error {norm_diff} and numerical error {numerical_error}"
+            )
+        return Y
+    except RuntimeError as e:
+        raise ValueError(
+            "Invalid values encountered during matrix multiplication: "
+            f"Input matrix: {mat.toarray()}"
+            f"system error: {e}"
+        ) from e
+
+
 def _compute_relative_entropy(rho, sigma, base=None):
     r"""
     Compute the quantum relative entropy of density matrix rho with respect to sigma.


Test patch:
diff --git a/tests/math/test_matrix_manipulation.py b/tests/math/test_matrix_manipulation.py
index 8ea97f2ee..9fe6c7bed 100644
--- a/tests/math/test_matrix_manipulation.py
+++ b/tests/math/test_matrix_manipulation.py
@@ -19,10 +19,12 @@ import numpy as np
 import pytest
 from gate_data import CNOT, II, SWAP, I, Toffoli
 from scipy.sparse import csr_matrix
+from scipy.stats import unitary_group
 
 import pennylane as qml
 from pennylane import numpy as pnp
-from pennylane.math import expand_matrix, expand_vector
+from pennylane.math import expand_matrix, expand_vector, sqrt_matrix, sqrt_matrix_sparse
+from pennylane.math.quantum import _denman_beavers_iterations
 
 # Define a list of dtypes to test
 dtypes = ["complex64", "complex128"]
@@ -1068,3 +1070,198 @@ class TestExpandVector:
         """Test exception raised if incorrect sized vector provided."""
         with pytest.raises(ValueError, match="Vector parameter must be of length"):
             expand_vector(TestExpandVector.VECTOR1, [0, 1], 4)
+
+
+class TestSqrtMatrix:
+    """Tests for the sqrt_matrix function."""
+
+    # NOTE: make sure the matrix is positive definite
+    dm_list = [
+        np.array([[1, 0], [0, 1]]),
+        np.array([[1, 0], [0, 2]]),
+        np.array([[4, 2], [2, 3]]),
+    ]
+    shape_list = range(2, 10)
+    # NOTE: sqrt_matrix is frequently used by BlockEncode
+    #       here below are some test matrices that are used
+    #       in the BlockEncode tests
+    matrices = [
+        # 2x2 matrices
+        np.array([[0.1, 0.2], [0.3, 0.4]]),
+        # Non-square matrices
+        np.array([[0.1, 0.2, 0.3], [0.3, 0.4, 0.2]]),
+        # 3x3 matrix
+        np.array([[0.1, 0.2, 0.3], [0.3, 0.4, 0.2], [0.1, 0.2, 0.3]]),
+        # Identity-like matrices
+        np.array([[1, 0], [0, 1]]),
+        np.identity(3),
+        # Matrix with complex entries
+        0.2 * np.array([[0.3, 0.9539392j], [0.9539392j, -0.3]]),
+    ]
+    # negative matrices: matrices that have negative eigenvalues
+    matrices_negative = [
+        # 2x2 matrix with negative eigenvalue
+        np.array([[1, 2], [2, -3]]),
+        # 3x3 matrix with mixed positive/negative eigenvalues
+        np.array([[2, -1, 0], [-1, -2, 1], [0, 1, 3]]),
+        # 4x4 matrix with negative eigenvalues
+        np.array([[1, 2, 0, 1], [2, -2, 1, 0], [0, 1, -3, 2], [1, 0, 2, -1]]),
+    ]
+    # Known problematic matrices with informative error message.
+    # NOTE: keep the entries here common. No fine-tuned corner cases.
+    illmats_info_pairs = [
+        (
+            np.array([[-1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]),
+            ValueError,
+            "Factor is exactly singular",
+        )
+    ]
+
+    @pytest.mark.parametrize("dm", dm_list)
+    def test_sqrt_matrix_sparse_dm(self, dm, tol):
+        """Test the sqrt_matrix function."""
+
+        A = qml.math.asarray(dm)
+        A_sparse = csr_matrix(A)
+
+        assert np.allclose(sqrt_matrix(A), sqrt_matrix_sparse(A_sparse).toarray(), atol=tol, rtol=0)
+
+    @pytest.mark.parametrize("shape", shape_list)
+    def test_sqrt_matrix_sparse_random(self, shape, tol):
+        """Test the sqrt_matrix function."""
+
+        # From unitary group
+        dm = unitary_group.rvs(shape)
+        dm = dm @ dm.T.conj()
+        A = qml.math.asarray(dm)
+        A_sparse = csr_matrix(A)
+
+        assert np.allclose(sqrt_matrix(A), sqrt_matrix_sparse(A_sparse).toarray(), atol=tol, rtol=0)
+
+    @pytest.mark.parametrize("matrix", matrices)
+    def test_sqrt_matrix_inputs(self, matrix, tol):
+        """Test sqrt_matrix function with various input matrices."""
+
+        # type M^† M
+        A = matrix.T.conj() @ matrix
+        A = np.eye(A.shape[0]) - A
+        A_sparse = csr_matrix(A)
+
+        result_sparse = sqrt_matrix_sparse(A_sparse)
+        result = result_sparse.toarray()
+        result_2 = result @ result
+
+        assert np.allclose(result_2, A, atol=tol, rtol=0)
+
+        # type M M^†
+        A = matrix @ matrix.T.conj()
+        A = np.eye(A.shape[0]) - A
+        A_sparse = csr_matrix(A)
+
+        result_sparse = sqrt_matrix_sparse(A_sparse)
+        result = result_sparse.toarray()
+        result_2 = result @ result
+
+        assert np.allclose(result_2, A, atol=tol, rtol=0)
+
+    @pytest.mark.parametrize("m", matrices_negative)
+    def test_sqrt_matrix_sparse_input_negative(self, m):
+        """Test that sqrt_matrix_sparse raises ValueError for matrices with negative eigenvalues."""
+        m = csr_matrix(m)
+        with pytest.raises(ValueError):
+            sqrt_matrix_sparse(m)
+
+    @pytest.mark.parametrize("m, e, info", illmats_info_pairs)
+    def test_sqrt_matrix_sparse_input_ill_conditioned(self, m, e, info):
+        """Test that appropriate errors are raised for ill-conditioned matrices"""
+        m = csr_matrix(m)
+        with pytest.raises(e, match=info):
+            sqrt_matrix_sparse(m)
+
+    def test_sqrt_matrix_sparse_input_valid(self):
+        """Test that if dense input errors raised"""
+        A = np.array([[1, 0], [0, 1]])
+        with pytest.raises(TypeError, match="only supports scipy.sparse matrices"):
+            sqrt_matrix_sparse(A)
+
+
+CONVERGENCE_ERROR = "Convergence threshold not reached"
+
+
+class TestDenmanBeaversIterations:
+    """Tests for the Denman-Beavers iteration method for matrix square root"""
+
+    def test_singular_matrix(self):
+        """Test that singular matrix raises appropriate error"""
+        n = 4
+        mat = csr_matrix(np.diag([0.0] + [1.0] * (n - 1)))
+
+        with pytest.raises(ValueError, match="Factor is exactly singular"):
+            _denman_beavers_iterations(mat)
+
+    def test_overflow_matrix(self):
+        """Test that matrix with very large values raises convergence warning"""
+        n = 4
+        mat = csr_matrix(np.eye(n) * 1e150)
+
+        with pytest.raises(ValueError, match=CONVERGENCE_ERROR):
+            _denman_beavers_iterations(mat)
+
+    def test_invalid_value_matrix(self):
+        """Test that matrix leading to invalid values raises appropriate error"""
+        n = 4
+        mat = csr_matrix(np.diag([-1e-10] + [1.0] * (n - 1)))
+
+        with pytest.raises(ValueError, match="nan or inf"):
+            _denman_beavers_iterations(mat)
+
+    def test_non_convergent_matrix(self):
+        """Test that non-convergent matrix raises appropriate error"""
+        mat = csr_matrix([[1, 1e-6], [1e-6, 1e-12 + 1e-25]])
+        with pytest.raises(ValueError, match=CONVERGENCE_ERROR):
+            _denman_beavers_iterations(mat)
+
+    def test_unstable_matrix(self):
+        """Test that numerically unstable matrix raises convergence warning"""
+        n = 4
+        mat = csr_matrix(np.diag([1e-200, 1e200] + [1.0] * (n - 2)))
+
+        with pytest.raises(ValueError, match=CONVERGENCE_ERROR):
+            _denman_beavers_iterations(mat)
+
+    @pytest.mark.parametrize("size", [2, 3, 4, 5])
+    def test_valid_positive_definite(self, size):
+        """Test that valid real, positive definite matrices work correctly"""
+        # Create a positive definite matrix
+        A = np.random.random((size, size))
+        mat = csr_matrix(np.eye(size) - 0.1 * (A @ A.T))
+
+        result = _denman_beavers_iterations(mat)
+        # Check that result is a valid square root
+        result_dense = result.toarray()
+        original_dense = mat.toarray()
+        qml.math.allclose(result_dense @ result_dense, original_dense, atol=1e-7, rtol=1e-7)
+
+    def test_hermitian_matrix(self):
+        """Test that Hermitian matrices work correctly on Hermitians of positive det.
+        Emulate random users' random input; the iteration should pass the simple branch determined by determinant, or it needs extra examine.
+        """
+        n = 4
+        A = np.random.random((n, n)) + 1j * np.random.random((n, n))
+        mat = np.eye(n) - 0.2 * (A @ A.T.conj())
+
+        def _reverse_det_hermitian(mat):
+            eigvals, eigvecs = np.linalg.eigh(mat)
+            eigvals[0] *= -1
+            return eigvecs @ np.diag(eigvals) @ eigvecs.T.conj()
+
+        det_mat = np.real(np.linalg.det(mat))
+        good_mat = mat if det_mat > 0 else _reverse_det_hermitian(mat)
+
+        result = _denman_beavers_iterations(csr_matrix(good_mat))
+        result_dense = result.toarray()
+        qml.math.allclose(result_dense @ result_dense, good_mat, atol=1e-7, rtol=1e-7)
+
+        bad_mat = mat if det_mat < 0 else _reverse_det_hermitian(mat)
+        with pytest.raises(ValueError, match="Invalid values encountered"):
+            _denman_beavers_iterations(csr_matrix(bad_mat))