add psd option for sqrtm

Author: refraction-ray

CHANGELOG.md

@@ -12,6 +12,8 @@
 
 - Add `u1_inds`, `u1_mask`, `u1_project`, and `u1_enlarge` functions in quantum.py as utilities in charged conservation systems
 
+- Add `psd` boolean to `sqrtmh` method for backends
+
 ### Fixed
 
 - Fix customized jax eigh operator by noting the return is a namedtuple

tensorcircuit/backends/abstract_backend.py

@@ -46,17 +46,22 @@ def expm(self: Any, a: Tensor) -> Tensor:
             "Backend '{}' has not implemented `expm`.".format(self.name)
         )
 
-    def sqrtmh(self: Any, a: Tensor) -> Tensor:
+    def sqrtmh(self: Any, a: Tensor, psd: bool = False) -> Tensor:
         """
         Return the sqrtm of a Hermitian matrix ``a``.
 
         :param a: tensor in matrix form
         :type a: Tensor
+        :param psd: whether the input ``a`` is guaranteed as a positive semidefinite matrix,
+            defaults False
+        :type psd: bool
         :return: sqrtm of ``a``
         :rtype: Tensor
         """
         # maybe friendly for AD and also cosidering that several backend has no support for native sqrtm
         e, v = self.eigh(a)
+        if psd:
+            e = self.relu(e)
         e = self.sqrt(e)
         return v @ self.diagflat(e) @ self.adjoint(v)
 

tests/test_backends.py

@@ -119,8 +119,12 @@ def test_backend_methods(backend):
     ans = np.array([[1, 0.5j], [-0.5j, 1]])
     ans2 = ans @ ans
     ansp = tc.backend.sqrtmh(tc.array_to_tensor(ans2))
-    print(ansp @ ansp, ans @ ans)
+    # print(ansp @ ansp, ans @ ans)
     np.testing.assert_allclose(ansp @ ansp, ans @ ans, atol=1e-4)
+    singularm = np.array([[4.0, 0], [0, -1e-3]])
+    np.testing.assert_allclose(
+        tc.backend.sqrtmh(singularm, psd=True), np.array([[2.0, 0], [0, 0]]), atol=1e-5
+    )
 
     np.testing.assert_allclose(
         tc.backend.sum(tc.array_to_tensor(np.arange(4))), 6, atol=1e-4