[ITensors] Customisable Arrow directions in SVD

src/tensor_operations/matrix_algebra.jl

@@ -87,3 +87,16 @@ 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)

src/tensor_operations/matrix_decomposition.jl

@@ -106,7 +106,13 @@ 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::Union{Nothing,ITensors.Arrow}=nothing,
+  rightdir::Union{Nothing,ITensors.Arrow}=nothing,
+  kwargs...,
+)
   utags::TagSet = get(kwargs, :lefttags, get(kwargs, :utags, "Link,u"))
   vtags::TagSet = get(kwargs, :righttags, get(kwargs, :vtags, "Link,v"))
 
@@ -133,11 +139,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,10 +535,32 @@ function factorize_qr(A::ITensor, Linds...; kwargs...)
   return L, R
 end
 
-function factorize_svd(A::ITensor, Linds...; kwargs...)
-  ortho::String = get(kwargs, :ortho, "left")
-  alg::String = get(kwargs, :svd_alg, "divide_and_conquer")
-  USV = svd(A, Linds...; kwargs..., alg=alg)
+#
+# 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)
+  (storage(D) isa Union{Diag,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,
+  ortho="left",
+  svd_alg="divide_and_conquer",
+  dir=ITensors.In,
+  kwargs...,
+)
+  leftdir, rightdir = -dir, -dir
+  USV = svd(A, Linds...; leftdir, rightdir, alg=svd_alg, kwargs...)
   if isnothing(USV)
     return nothing
   end
@@ -544,14 +570,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 = U * sqrtDL, 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
 
@@ -646,13 +674,6 @@ function factorize(A::ITensor, Linds...; kwargs...)
     which_decomp = "eigen"
   end
 
-  # Deprecated keywords
-  if haskey(kwargs, :dir)
-    error("""dir keyword in factorize has been replace by ortho.
-    Note that the default is now `left`, meaning for the results
-    L,R = factorize(A), L forms an orthogonal basis.""")
-  end
-
   if haskey(kwargs, :which_factorization)
     error("""which_factorization keyword in factorize has
     been replace by which_decomp.""")

test/base/test_decomp.jl

@@ -111,7 +111,6 @@ end
     i = Index(2, "i")
     j = Index(2, "j")
     A = randomITensor(i, j)
-    @test_throws ErrorException factorize(A, i; dir="left")
     @test_throws ErrorException factorize(A, i; ortho="fakedir")
   end
 

test/base/test_fermions.jl

@@ -744,6 +744,21 @@ import ITensors: Out, In
       VV = dag(V) * prime(V, v)
       @test norm(VV - id(v)) ≈ 0
     end
+
+    #Factorize SVD Test. Specifying arrows on S.
+    let
+      l1, l2 = Index(QN("Nf", -1) => 1, QN("Nf", 1) => 1; tags="l1", dir=ITensors.In),
+      Index(QN("Nf", 2) => 1, QN("Nf", 1) => 1; tags="l2", dir=ITensors.Out)
+      r1, r2, r3 = Index(QN("Nf", -2) => 1, QN("Nf", 1) => 1; tags="r1", dir=ITensors.Out),
+      Index(QN("Nf", 2) => 1, QN("Nf", 1) => 1; tags="r2", dir=ITensors.In),
+      Index(QN("Nf", -2) => 1, QN("Nf", 1) => 1; tags="r3", dir=ITensors.In)
+      A = randomITensor(l1, l2, r1, r2, r3)
+
+      for dir in [ITensors.Out, ITensors.In]
+        L, R, spec = ITensors.factorize_svd(A, l1, l2; dir, ortho="none")
+        @test norm(L * R - A) <= 1e-14
+      end
+    end
   end
 
   @testset "Fermion Contraction with Combined Indices" begin

test/base/test_svd.jl

@@ -186,6 +186,31 @@ include(joinpath(@__DIR__, "utils", "util.jl"))
     @test Base.eltype(V) === eltype
   end
 
+  @testset "svd arrow directions" begin
+    l1, l2 = Index(QN("Sz", -1) => 1, QN("Sz", 1) => 1; tags="l1", dir=ITensors.In),
+    Index(QN("Sz", 2) => 1, QN("Sz", 1) => 1; tags="l2", dir=ITensors.Out)
+    r1, r2, r3 = Index(QN("Sz", -2) => 1, QN("Sz", 1) => 1; tags="r1", dir=ITensors.Out),
+    Index(QN("Sz", 2) => 1, QN("Sz", 1) => 1; tags="r2", dir=ITensors.In),
+    Index(QN("Sz", -2) => 1, QN("Sz", 1) => 1; tags="r3", dir=ITensors.In)
+    A = randomITensor(l1, l2, r1, r2, r3)
+
+    for leftdir in [ITensors.Out, ITensors.In]
+      for rightdir in [ITensors.Out, ITensors.In]
+        U, S, V = svd(A, l1, l2; leftdir, rightdir)
+        s1, s2 = inds(S)
+        @test dir(s1) == leftdir
+        @test dir(s2) == rightdir
+        @test norm(U * S * V - A) <= 1e-14
+      end
+    end
+
+    for dir in [ITensors.Out, ITensors.In]
+      L, R, spec = ITensors.factorize_svd(A, l1, l2; dir, ortho="none")
+      @test dir == ITensors.dir(commonind(L, R))
+      @test norm(L * R - A) <= 1e-14
+    end
+  end
+
   # TODO: remove this test, it takes a long time
   ## @testset "Ill-conditioned matrix" begin
   ##   d = 5000