Add absorb_s option for SVD (#9)

* Add absorb_s option for svd, einsvd, einsumsvd

And __ipow__ is now supported via __pow__ by python

* Add test

tensorbackends/backends/ctf/ctf_backend.py

@@ -7,6 +7,7 @@
 
 from ...interface import Backend
 from ...utils import einstr
+from ...utils.svd_absorb_s import svd_absorb_s, svd_absorb_s_ctf
 from .ctf_random import CTFRandom
 from .ctf_tensor import CTFTensor
 
@@ -82,35 +83,35 @@ def einsum(self, subscripts, *operands):
         expr = einstr.parse_einsum(subscripts, ndims)
         return self._einsum(expr, operands)
 
-    def einsvd_reduced(self, subscripts, a, rank=None):
+    def einsvd_reduced(self, subscripts, a, rank=None, absorb_s=False):
         if not isinstance(a, self.tensor):
             raise TypeError('the input should be {}'.format(self.tensor.__qualname__))
         expr = einstr.parse_einsvd(subscripts, a.ndim)
-        return self._einsvd_reduced(expr, a, rank)
+        return self._einsvd_reduced(expr, a, rank, absorb_s)
 
-    def einsvd_rand(self, subscripts, a, rank, niter=1, oversamp=5):
+    def einsvd_rand(self, subscripts, a, rank, niter=1, oversamp=5, absorb_s=False):
         if not isinstance(a, self.tensor):
             raise TypeError('the input should be {}'.format(self.tensor.__qualname__))
         expr = einstr.parse_einsvd(subscripts, a.ndim)
-        return self._einsvd_rand(expr, a, rank, niter, oversamp)
+        return self._einsvd_rand(expr, a, rank, niter, oversamp, absorb_s=absorb_s)
 
-    def einsumsvd_reduced(self, subscripts, *operands, rank=None):
+    def einsumsvd_reduced(self, subscripts, *operands, rank=None, absorb_s=False):
         if not all(isinstance(operand, self.tensor) for operand in operands):
             raise TypeError('all operands should be {}'.format(self.tensor.__qualname__))
         ndims = [operand.ndim for operand in operands]
         expr = einstr.parse_einsumsvd(subscripts, ndims)
         einsum_expr, einsvd_expr = einstr.split_einsumsvd(expr)
         a = self._einsum(einsum_expr, operands)
-        return self._einsvd_reduced(einsvd_expr, a, rank)
+        return self._einsvd_reduced(einsvd_expr, a, rank, absorb_s=absorb_s)
 
-    def einsumsvd_rand(self, subscripts, *operands, rank, niter=1, oversamp=5):
+    def einsumsvd_rand(self, subscripts, *operands, rank, niter=1, oversamp=5, absorb_s=False):
         if not all(isinstance(operand, self.tensor) for operand in operands):
             raise TypeError('all operands should be {}'.format(self.tensor.__qualname__))
         ndims = [operand.ndim for operand in operands]
         expr = einstr.parse_einsumsvd(subscripts, ndims)
         einsum_expr, einsvd_expr = einstr.split_einsumsvd(expr)
         a = self._einsum(einsum_expr, operands)
-        return self._einsvd_rand(einsvd_expr, a, rank, niter, oversamp)
+        return self._einsvd_rand(einsvd_expr, a, rank, niter, oversamp, absorb_s=absorb_s)
 
     def isclose(self, a, b, *, rtol=1e-9, atol=0.0):
         return abs(a - b) <= atol + rtol * abs(b)
@@ -122,13 +123,15 @@ def inv(self, a):
         u, s, v = self.einsvd('ij->ia,ja', a)
         return self.einsum('ia,a,ja->ji', u, 1/s, v)
 
-    def svd(self, a):
+    def svd(self, a, absorb_s=False):
         if not isinstance(a, self.tensor):
             raise TypeError('the input should be {}'.format(self.tensor.__qualname__))
         if a.ndim != 2:
             raise TypeError('the input tensor should be a matrix')
         u, s, vh = ctf.svd(a.unwrap())
-        return self.tensor(u), self.tensor(ctf.real(s)), self.tensor(vh)
+        u, s, vh = self.tensor(u), self.tensor(ctf.real(s)), self.tensor(vh)
+        u, s, vh = svd_absorb_s(u, s, vh, absorb_s)
+        return u, s, vh
 
     def __getattr__(self, attr):
         wrap = lambda val: CTFTensor(val) if isinstance(val, ctf.tensor) else val
@@ -167,19 +170,18 @@ def _einsum(self, expr, operands):
         else:
             return result
 
-    def _einsvd_reduced(self, expr, a, rank):
-        u, s, vh = a.tsr.i(expr.inputs[0].indices_string).svd(
-            expr.outputs[0].indices_string,
-            expr.outputs[1].indices_string,
-            rank=rank,
-        )
+    def _einsvd_reduced(self, expr, a, rank, absorb_s):
+        u_str, vh_str = expr.outputs[0].indices_string, expr.outputs[1].indices_string
+        u, s, vh = a.tsr.i(expr.inputs[0].indices_string).svd(u_str, vh_str, rank=rank)
+        u, s, vh = self.tensor(u), self.tensor(ctf.real(s)), self.tensor(vh)
+        u, s, vh = svd_absorb_s_ctf(u, s, vh, absorb_s, u_str, vh_str)
         u_newshape = expr.outputs[0].newshape(u.shape)
         vh_newshape = expr.outputs[1].newshape(vh.shape)
         if u_newshape != u.shape: u = u.reshape(*u_newshape)
         if vh_newshape != vh.shape: vh = vh.reshape(*vh_newshape)
-        return self.tensor(u), self.tensor(ctf.real(s)), self.tensor(vh)
+        return u, s, vh
 
-    def _einsvd_rand(self, expr, a, rank, niter, oversamp):
+    def _einsvd_rand(self, expr, a, rank, niter, oversamp, absorb_s):
         newindex = (expr.output_indices - expr.input_indices).pop()
         axis_of_index = {index: axis for axis, index in enumerate(expr.inputs[0])}
         u_axes_from_a = [axis_of_index[index] for index in expr.outputs[0] if index != newindex]
@@ -188,7 +190,7 @@ def _einsvd_rand(self, expr, a, rank, niter, oversamp):
         a_matrix_axes = [*u_axes_from_a, *vh_axes_from_a]
         a_matrix_shape = (np.prod([a.shape[axis] for axis in u_axes_from_a]), -1)
         a_matrix = a.transpose(*a_matrix_axes).reshape(*a_matrix_shape)
-        u, s, vh = self.rsvd(a_matrix, rank, niter, oversamp)
+        u, s, vh = self.rsvd(a_matrix, rank, niter, oversamp, absorb_s)
         # form u
         u = u.reshape(*(a.shape[axis] for axis in u_axes_from_a), s.shape[0])
         u = self.moveaxis(u, -1, expr.outputs[0].find(newindex))

tensorbackends/backends/ctf/ctf_tensor.py

@@ -145,7 +145,7 @@ def method(self, other):
     '__imatmul__',
     '__itruediv__',
     '__ifloordiv__',
-    '__ipow__',
+    # '__ipow__', # TODO Add __ipow__ to ctf tensor
 
     '__lt__',
     '__le__',

tensorbackends/backends/ctfview/ctfview_backend.py

@@ -7,6 +7,7 @@
 
 from ...interface import Backend
 from ...utils import einstr
+from ...utils.svd_absorb_s import svd_absorb_s
 from .ctfview_random import CTFViewRandom
 from .ctfview_tensor import CTFViewTensor
 from . import indices_utils
@@ -91,48 +92,48 @@ def einsum(self, subscripts, *operands):
         expr = einstr.parse_einsum(subscripts, ndims)
         return self._einsum(expr, operands)
 
-    def einsvd_reduced(self, subscripts, a, rank=None):
+    def einsvd_reduced(self, subscripts, a, rank=None, absorb_s=False):
         if not isinstance(a, self.tensor):
             raise TypeError('the input should be {}'.format(self.tensor.__qualname__))
         expr = einstr.parse_einsvd(subscripts, a.ndim)
         def svd_func(matrix):
-            u, s, vh = self.svd(matrix)
+            u, s, vh = self.svd(matrix, absorb_s=absorb_s)
             if rank is not None and s.shape[0] > rank:
                 u, s, vh = u[:,:rank], s[:rank], vh[:rank,:]
             return u, s, vh
         return self._einsvd(expr, a, svd_func)
 
-    def einsvd_rand(self, subscripts, a, rank, niter=1, oversamp=5):
+    def einsvd_rand(self, subscripts, a, rank, niter=1, oversamp=5, absorb_s=False):
         if not isinstance(a, self.tensor):
             raise TypeError('the input should be {}'.format(self.tensor.__qualname__))
         expr = einstr.parse_einsvd(subscripts, a.ndim)
         def svd_func(matrix):
-            return self.rsvd(matrix, rank, niter, oversamp)
+            return self.rsvd(matrix, rank, niter, oversamp, absorb_s=absorb_s)
         return self._einsvd(expr, a, svd_func)
 
-    def einsumsvd_reduced(self, subscripts, *operands, rank=None):
+    def einsumsvd_reduced(self, subscripts, *operands, rank=None, absorb_s=False):
         if not all(isinstance(operand, self.tensor) for operand in operands):
             raise TypeError('all operands should be {}'.format(self.tensor.__qualname__))
         ndims = [operand.ndim for operand in operands]
         expr = einstr.parse_einsumsvd(subscripts, ndims)
         einsum_expr, einsvd_expr = einstr.split_einsumsvd(expr)
         a = self._einsum(einsum_expr, operands)
         def svd_func(matrix):
-            u, s, vh = self.svd(matrix)
+            u, s, vh = self.svd(matrix, absorb_s=absorb_s)
             if rank is not None and s.shape[0] > rank:
                 u, s, vh = u[:,:rank], s[:rank], vh[:rank,:]
             return u, s, vh
         return self._einsvd(einsvd_expr, a, svd_func)
 
-    def einsumsvd_rand(self, subscripts, *operands, rank, niter=1, oversamp=5):
+    def einsumsvd_rand(self, subscripts, *operands, rank, niter=1, oversamp=5, absorb_s=False):
         if not all(isinstance(operand, self.tensor) for operand in operands):
             raise TypeError('all operands should be {}'.format(self.tensor.__qualname__))
         ndims = [operand.ndim for operand in operands]
         expr = einstr.parse_einsumsvd(subscripts, ndims)
         einsum_expr, einsvd_expr = einstr.split_einsumsvd(expr)
         a = self._einsum(einsum_expr, operands)
         def svd_func(matrix):
-            return self.rsvd(matrix, rank, niter, oversamp)
+            return self.rsvd(matrix, rank, niter, oversamp, absorb_s=absorb_s)
         return self._einsvd(einsvd_expr, a, svd_func)
 
     def isclose(self, a, b, *, rtol=1e-9, atol=0.0):
@@ -147,13 +148,15 @@ def inv(self, a):
         u, s, v = self.einsvd('ij->ia,ja', a)
         return self.einsum('ia,a,ja->ji', u, 1/s, v)
 
-    def svd(self, a):
+    def svd(self, a, absorb_s=False):
         if not isinstance(a, self.tensor):
             raise TypeError('the input should be {}'.format(self.tensor.__qualname__))
         if a.ndim != 2:
             raise TypeError('the input tensor should be a matrix')
         u, s, vh = ctf.svd(a.unwrap())
-        return self.tensor(u), self.tensor(ctf.real(s)), self.tensor(vh)
+        u, s, vh = self.tensor(u), self.tensor(ctf.real(s)), self.tensor(vh)
+        u, s, vh = svd_absorb_s(u, s, vh, absorb_s)
+        return u, s, vh
 
     def __getattr__(self, attr):
         wrap = lambda val: CTFViewTensor(val) if isinstance(val, ctf.tensor) else val

tensorbackends/backends/ctfview/ctfview_tensor.py

@@ -200,7 +200,7 @@ def method(self, other):
     '__imatmul__',
     '__itruediv__',
     '__ifloordiv__',
-    '__ipow__',
+    # '__ipow__', # TODO Add __ipow__ to ctf tensor
 
     '__lt__',
     '__le__',

tensorbackends/backends/cupy/cupy_backend.py

@@ -9,6 +9,7 @@
 
 from ...interface import Backend
 from ...utils import einstr
+from ...utils.svd_absorb_s import svd_absorb_s
 from .cupy_random import CuPyRandom
 from .cupy_tensor import CuPyTensor
 
@@ -84,48 +85,48 @@ def einsum(self, subscripts, *operands):
         expr = einstr.parse_einsum(subscripts, ndims)
         return self._einsum(expr, operands)
 
-    def einsvd_reduced(self, subscripts, a, rank=None):
+    def einsvd_reduced(self, subscripts, a, rank=None, absorb_s=False):
         if not isinstance(a, self.tensor):
             raise TypeError('the input should be {}'.format(self.tensor.__qualname__))
         expr = einstr.parse_einsvd(subscripts, a.ndim)
         def svd_func(matrix):
-            u, s, vh = self.svd(matrix)
+            u, s, vh = self.svd(matrix, absorb_s=absorb_s)
             if rank is not None and s.shape[0] > rank:
                 u, s, vh = u[:,:rank], s[:rank], vh[:rank,:]
             return u, s, vh
         return self._einsvd(expr, a, svd_func)
 
-    def einsvd_rand(self, subscripts, a, rank, niter=1, oversamp=5):
+    def einsvd_rand(self, subscripts, a, rank, niter=1, oversamp=5, absorb_s=False):
         if not isinstance(a, self.tensor):
             raise TypeError('the input should be {}'.format(self.tensor.__qualname__))
         expr = einstr.parse_einsvd(subscripts, a.ndim)
         def svd_func(matrix):
-            return self.rsvd(matrix, rank, niter, oversamp)
+            return self.rsvd(matrix, rank, niter, oversamp, absorb_s=absorb_s)
         return self._einsvd(expr, a, svd_func)
 
-    def einsumsvd_reduced(self, subscripts, *operands, rank=None):
+    def einsumsvd_reduced(self, subscripts, *operands, rank=None, absorb_s=False):
         if not all(isinstance(operand, self.tensor) for operand in operands):
             raise TypeError('all operands should be {}'.format(self.tensor.__qualname__))
         ndims = [operand.ndim for operand in operands]
         expr = einstr.parse_einsumsvd(subscripts, ndims)
         einsum_expr, einsvd_expr = einstr.split_einsumsvd(expr)
         a = self._einsum(einsum_expr, operands)
         def svd_func(matrix):
-            u, s, vh = self.svd(matrix)
+            u, s, vh = self.svd(matrix, absorb_s=absorb_s)
             if rank is not None and s.shape[0] > rank:
                 u, s, vh = u[:,:rank], s[:rank], vh[:rank,:]
             return u, s, vh
         return self._einsvd(einsvd_expr, a, svd_func)
 
-    def einsumsvd_rand(self, subscripts, *operands, rank, niter=1, oversamp=5):
+    def einsumsvd_rand(self, subscripts, *operands, rank, niter=1, oversamp=5, absorb_s=False):
         if not all(isinstance(operand, self.tensor) for operand in operands):
             raise TypeError('all operands should be {}'.format(self.tensor.__qualname__))
         ndims = [operand.ndim for operand in operands]
         expr = einstr.parse_einsumsvd(subscripts, ndims)
         einsum_expr, einsvd_expr = einstr.split_einsumsvd(expr)
         a = self._einsum(einsum_expr, operands)
         def svd_func(matrix):
-            return self.rsvd(matrix, rank, niter, oversamp)
+            return self.rsvd(matrix, rank, niter, oversamp, absorb_s=absorb_s)
         return self._einsvd(einsvd_expr, a, svd_func)
 
     def isclose(self, a, b, *, rtol=1e-9, atol=0.0):
@@ -142,9 +143,11 @@ def allclose(self, a, b, *, rtol=1e-9, atol=0.0):
     def inv(self, a):
         return CuPyTensor(la.inv(a.unwrap()))
 
-    def svd(self, a):
+    def svd(self, a, absorb_s=False):
         u, s, vh = la.svd(a.unwrap(), full_matrices=False)
-        return CuPyTensor(u), CuPyTensor(s), CuPyTensor(vh)
+        u, s, vh = self.tensor(u), self.tensor(s), self.tensor(vh)
+        u, s, vh = svd_absorb_s(u, s, vh, absorb_s)
+        return u, s, vh
 
     def __getattr__(self, attr):
         wrap = lambda val: CuPyTensor(val) if isinstance(val, cp.ndarray) else val

tensorbackends/backends/numpy/numpy_backend.py

@@ -10,6 +10,7 @@
 
 from ...interface import Backend
 from ...utils import einstr
+from ...utils.svd_absorb_s import svd_absorb_s
 from .numpy_random import NumPyRandom
 from .numpy_tensor import NumPyTensor
 
@@ -85,48 +86,48 @@ def einsum(self, subscripts, *operands):
         expr = einstr.parse_einsum(subscripts, ndims)
         return self._einsum(expr, operands)
 
-    def einsvd_reduced(self, subscripts, a, rank=None):
+    def einsvd_reduced(self, subscripts, a, rank=None, absorb_s=False):
         if not isinstance(a, self.tensor):
             raise TypeError('the input should be {}'.format(self.tensor.__qualname__))
         expr = einstr.parse_einsvd(subscripts, a.ndim)
         def svd_func(matrix):
-            u, s, vh = self.svd(matrix)
+            u, s, vh = self.svd(matrix, absorb_s=absorb_s)
             if rank is not None and s.shape[0] > rank:
                 u, s, vh = u[:,:rank], s[:rank], vh[:rank,:]
             return u, s, vh
         return self._einsvd(expr, a, svd_func)
 
-    def einsvd_rand(self, subscripts, a, rank, niter=1, oversamp=5):
+    def einsvd_rand(self, subscripts, a, rank, niter=1, oversamp=5, absorb_s=False):
         if not isinstance(a, self.tensor):
             raise TypeError('the input should be {}'.format(self.tensor.__qualname__))
         expr = einstr.parse_einsvd(subscripts, a.ndim)
         def svd_func(matrix):
-            return self.rsvd(matrix, rank, niter, oversamp)
+            return self.rsvd(matrix, rank, niter, oversamp, absorb_s=absorb_s)
         return self._einsvd(expr, a, svd_func)
 
-    def einsumsvd_reduced(self, subscripts, *operands, rank=None):
+    def einsumsvd_reduced(self, subscripts, *operands, rank=None, absorb_s=False):
         if not all(isinstance(operand, self.tensor) for operand in operands):
             raise TypeError('all operands should be {}'.format(self.tensor.__qualname__))
         ndims = [operand.ndim for operand in operands]
         expr = einstr.parse_einsumsvd(subscripts, ndims)
         einsum_expr, einsvd_expr = einstr.split_einsumsvd(expr)
         a = self._einsum(einsum_expr, operands)
         def svd_func(matrix):
-            u, s, vh = self.svd(matrix)
+            u, s, vh = self.svd(matrix, absorb_s=absorb_s)
             if rank is not None and s.shape[0] > rank:
                 u, s, vh = u[:,:rank], s[:rank], vh[:rank,:]
             return u, s, vh
         return self._einsvd(einsvd_expr, a, svd_func)
 
-    def einsumsvd_rand(self, subscripts, *operands, rank, niter=1, oversamp=5):
+    def einsumsvd_rand(self, subscripts, *operands, rank, niter=1, oversamp=5, absorb_s=False):
         if not all(isinstance(operand, self.tensor) for operand in operands):
             raise TypeError('all operands should be {}'.format(self.tensor.__qualname__))
         ndims = [operand.ndim for operand in operands]
         expr = einstr.parse_einsumsvd(subscripts, ndims)
         einsum_expr, einsvd_expr = einstr.split_einsumsvd(expr)
         a = self._einsum(einsum_expr, operands)
         def svd_func(matrix):
-            return self.rsvd(matrix, rank, niter, oversamp)
+            return self.rsvd(matrix, rank, niter, oversamp, absorb_s=absorb_s)
         return self._einsvd(einsvd_expr, a, svd_func)
 
     def isclose(self, a, b, *, rtol=1e-9, atol=0.0):
@@ -143,9 +144,11 @@ def allclose(self, a, b, *, rtol=1e-9, atol=0.0):
     def inv(self, a):
         return NumPyTensor(la.inv(a.unwrap()))
 
-    def svd(self, a):
+    def svd(self, a, absorb_s=False):
         u, s, vh = la.svd(a.unwrap(), full_matrices=False)
-        return NumPyTensor(u), NumPyTensor(s), NumPyTensor(vh)
+        u, s, vh = self.tensor(u), self.tensor(s), self.tensor(vh)
+        u, s, vh = svd_absorb_s(u, s, vh, absorb_s)
+        return u, s, vh
 
     def __getattr__(self, attr):
         wrap = lambda val: NumPyTensor(val) if isinstance(val, np.ndarray) else val

tensorbackends/extensions/einsumsvd_implicit_rand.py

@@ -4,7 +4,7 @@
 import scipy.linalg as la
 
 
-def einsumsvd_implicit_rand(backend, subscripts, *operands, rank, niter, orth_method):
+def einsumsvd_implicit_rand(backend, subscripts, *operands, rank, niter, orth_method, absorb_s):
     if orth_method == 'qr':
         orthogonalize = orthogonalize_qr
     elif orth_method == 'local_gram':
@@ -74,7 +74,7 @@ def einsumsvd_implicit_rand(backend, subscripts, *operands, rank, niter, orth_me
     op_YT = orthogonalize(backend, op_YT)
 
     op_X = apply_A(backend,expr_A,ops_A,term_YT,op_YT,term_X)
-    mat_U, S, mat_XVT = backend.svd(op_X.reshape(np.prod(op_X.shape)//r, r))
+    mat_U, S, mat_XVT = backend.svd(op_X.reshape(np.prod(op_X.shape)//r, r), absorb_s=absorb_s)
     op_YT = backend.tensordot(op_YT.conj(), mat_XVT, axes=((-1),(-1)))
     op_X = mat_U.reshape(*op_X.shape)
     U = op_X

tensorbackends/extensions/rsvd.py

@@ -1,5 +1,6 @@
+from ..utils.svd_absorb_s import svd_absorb_s
 
-def rsvd(backend, a, rank, niter, oversamp):
+def rsvd(backend, a, rank, niter, oversamp, absorb_s):
     dtype = a.dtype
     m, n = a.shape
     r = min(rank + oversamp, m, n)
@@ -17,4 +18,5 @@ def rsvd(backend, a, rank, niter, oversamp):
     u = q @ u_sub
     if rank < r:
         u, s, vh = u[:,:rank], s[:rank], vh[:rank,:]
+    u, s, vh = svd_absorb_s(u, s, vh, absorb_s)
     return u, s, vh