[ITensors] Customisable Arrow directions in SVD

src/tensor_operations/matrix_algebra.jl

@@ -87,3 +87,25 @@ function exp(A::ITensor; kwargs...)
   Lis = Ris'
   return exp(A, Lis, Ris; kwargs...)
 end
+
+function map_diag!(f::Function, it_destination::ITensor, it_source::ITensor)
+  return itensor(map_diag!(f, tensor(it_destination), tensor(it_source)))
+end
+map_diag(f::Function, it::ITensor) = map_diag!(f, copy(it), it)
+
+function map_diag!(f::Function, t_destination::Tensor, t_source::Tensor)
+  for i in 1:diaglength(t_destination)
+    NDTensors.setdiagindex!(t_destination, f(NDTensors.getdiagindex(t_source, i)), i)
+  end
+  return t_destination
+end
+map_diag(f::Function, t::Tensor) = map_diag!(f, copy(t), t)
+
+# Convenience functions
+sqrt_diag(it::ITensor) = map_diag(sqrt, it)
+inv_diag(it::ITensor) = map_diag(inv, it)
+invsqrt_diag(it::ITensor) = map_diag(inv ∘ sqrt, it)
+pinv_diag(it::ITensor) = map_diag(pinv, it)
+pinvsqrt_diag(it::ITensor) = map_diag(pinv ∘ sqrt, it)
+sqrtabs_diag(it::ITensor) = map_diag(sqrt ∘ abs, it)
+inv_sqrtabs_diag(it::ITensor) = map_diag(inv ∘ sqrt ∘ abs, it)

src/tensor_operations/matrix_decomposition.jl

@@ -106,7 +106,7 @@ Utrunc2, Strunc2, Vtrunc2 = svd(A, i, k; cutoff=1e-10);
 
 See also: [`factorize`](@ref), [`eigen`](@ref)
 """
-function svd(A::ITensor, Linds...; kwargs...)
+function svd(A::ITensor, Linds...; leftdir=nothing, rightdir=nothing, kwargs...)
   utags::TagSet = get(kwargs, :lefttags, get(kwargs, :utags, "Link,u"))
   vtags::TagSet = get(kwargs, :righttags, get(kwargs, :vtags, "Link,v"))
 
@@ -133,11 +133,9 @@ function svd(A::ITensor, Linds...; kwargs...)
     Ris = [α]
   end
 
-  CL = combiner(Lis...)
-  CR = combiner(Ris...)
-
+  CL = combiner(Lis...; dir=leftdir)
+  CR = combiner(Ris...; dir=rightdir)
   AC = A * CR * CL
-
   cL = combinedind(CL)
   cR = combinedind(CR)
   if inds(AC) != (cL, cR)
@@ -531,7 +529,22 @@ function factorize_qr(A::ITensor, Linds...; kwargs...)
   return L, R
 end
 
-function factorize_svd(A::ITensor, Linds...; kwargs...)
+#
+# Generic function implementing a square root decomposition of a diagonal, order 2 tensor with inds u, v
+#
+function sqrt_decomp(D::ITensor, u::Index, v::Index)
+  (storagetype(D) <: Diag || storagetype(D) <: DiagBlockSparse) ||
+    error("Must be a diagonal matrix ITensor.")
+  sqrtDL = diagITensor(u, dag(u)')
+  sqrtDR = diagITensor(v, dag(v)')
+  map_diag!(sqrt ∘ abs, sqrtDL, D)
+  map_diag!(sqrt ∘ abs, sqrtDR, D)
+  δᵤᵥ = copy(D)
+  map_diag!(sign, δᵤᵥ, D)
+  return sqrtDL, prime(δᵤᵥ), sqrtDR
+end
+
+function factorize_svd(A::ITensor, Linds...; (singular_values!)=nothing, kwargs...)
   ortho::String = get(kwargs, :ortho, "left")
   alg::String = get(kwargs, :svd_alg, "divide_and_conquer")
   USV = svd(A, Linds...; kwargs..., alg=alg)
@@ -544,14 +557,16 @@ function factorize_svd(A::ITensor, Linds...; kwargs...)
   elseif ortho == "right"
     L, R = U * S, V
   elseif ortho == "none"
-    sqrtS = S
-    sqrtS .= sqrt.(S)
-    L, R = U * sqrtS, sqrtS * V
-    replaceind!(L, v, u)
+    sqrtDL, δᵤᵥ, sqrtDR = sqrt_decomp(S, u, v)
+    sqrtDR = denseblocks(sqrtDR) * denseblocks(δᵤᵥ)
+    L, R = noprime(U * sqrtDL), noprime(V * sqrtDR)
   else
     error("In factorize using svd decomposition, ortho keyword
     $ortho not supported. Supported options are left, right, or none.")
   end
+
+  !isnothing(singular_values!) && (singular_values![] = S)
+
   return L, R, spec
 end
 

test/base/test_svd.jl

@@ -186,6 +186,51 @@ include(joinpath(@__DIR__, "utils", "util.jl"))
     @test Base.eltype(V) === eltype
   end
 
+  @testset "svd arrow directions" begin
+
+    #Test Spins
+    n = 10
+    χ = 5
+    s = siteinds("S=1/2", n; conserve_qns=true)
+    state = [isodd(i) ? "Up" : "Dn" for i in 1:n]
+    ψ = randomMPS(s, state; linkdims=χ)
+    A = ψ[4] * ψ[5]
+    linds = [inds(A)[1], inds(A)[2]]
+
+    for leftdir in [ITensors.Out, ITensors.In]
+      for rightdir in [ITensors.Out, ITensors.In]
+        U, S, V = svd(A, linds...; leftdir, rightdir)
+        s1, s2 = inds(S)[1], inds(S)[2]
+        @test ITensors.dir(s1) == leftdir
+        @test ITensors.dir(s2) == rightdir
+        @test norm(U * S * V - A) <= 1e-14
+
+        L, R, spec = ITensors.factorize_svd(
+          A, linds...; leftdir=leftdir, rightdir=rightdir, ortho="none"
+        )
+        @test norm(L * R - A) <= 1e-14
+      end
+    end
+
+    #Test Fermions
+    #WARNING: CURRENTLY ONLY STABLE FOR BOTH ARROW DIRECTIONS = OUT ON THE SVD. 
+    ITensors.enable_auto_fermion()
+    n = 5
+    χ = 2
+    s = siteinds("Fermion", n; conserve_qns=true)
+    state = [isodd(i) ? "Occ" : "Emp" for i in 1:n]
+    ψ = randomMPS(s, state; linkdims=χ)
+    A = ψ[3] * ψ[4]
+    linds = [inds(A)[1]]
+
+    L, R, spec = ITensors.factorize_svd(
+      A, linds...; leftdir=ITensors.Out, rightdir=ITensors.Out, ortho="none"
+    )
+    @test norm(L * R - A) <= 1e-14
+
+    ITensors.disable_auto_fermion()
+  end
+
   # TODO: remove this test, it takes a long time
   ## @testset "Ill-conditioned matrix" begin
   ##   d = 5000