more careful about gaugefixing svd_trunc

Author: lkdvos

src/common/gauge.jl

@@ -12,53 +12,42 @@ is real and positive.
 function gaugefix!(::Union{typeof(eig_full!), typeof(eigh_full!), typeof(gen_eig_full!)}, V::AbstractMatrix)
     for j in axes(V, 2)
         v = view(V, :, j)
-        s = conj(sign(_argmaxabs(v)))
-        @inbounds v .*= s
+        s = sign(_argmaxabs(v))
+        @inbounds v .*= conj(s)
     end
     return V
 end
 
 function gaugefix!(::typeof(svd_full!), U, Vᴴ)
-    m, n = size(U, 1), size(Vᴴ, 2)
+    m, n = size(U, 2), size(Vᴴ, 1)
     for j in 1:max(m, n)
         if j <= min(m, n)
             u = view(U, :, j)
             v = view(Vᴴ, j, :)
-            s = conj(sign(_argmaxabs(u)))
-            u .*= s
-            v .*= conj(s)
+            s = sign(_argmaxabs(u))
+            u .*= conj(s)
+            v .*= s
         elseif j <= m
             u = view(U, :, j)
-            s = conj(sign(_argmaxabs(u)))
-            u .*= s
+            s = sign(_argmaxabs(u))
+            u .*= conj(s)
         else
             v = view(Vᴴ, j, :)
-            s = conj(sign(_argmaxabs(v)))
-            v .*= s
+            s = sign(_argmaxabs(v))
+            v .*= conj(s)
         end
     end
     return (U, Vᴴ)
 end
 
-function gaugefix!(::typeof(svd_compact!), U, Vᴴ)
-    for j in 1:size(U, 2)
-        u = view(U, :, j)
-        v = view(Vᴴ, j, :)
-        s = conj(sign(_argmaxabs(u)))
-        u .*= s
-        v .*= conj(s)
-    end
-    return (U, Vᴴ)
-end
-
-function gaugefix!(::typeof(svd_trunc!), U, Vᴴ)
-    m, n = size(U, 1), size(Vᴴ, 2)
-    for j in 1:min(m, n)
+function gaugefix!(::Union{typeof(svd_compact!), typeof(svd_trunc!)}, U, Vᴴ)
+    @assert axes(U, 2) == axes(Vᴴ, 1)
+    for j in axes(U, 2)
         u = view(U, :, j)
         v = view(Vᴴ, j, :)
-        s = conj(sign(_argmaxabs(u)))
-        u .*= s
-        v .*= conj(s)
+        s = sign(_argmaxabs(u))
+        u .*= conj(s)
+        v .*= s
     end
     return (U, Vᴴ)
 end

src/implementations/svd.jl

@@ -1,4 +1,4 @@
-# Inputs
+# Input
 # ------
 copy_input(::typeof(svd_full), A::AbstractMatrix) = copy!(similar(A, float(eltype(A))), A)
 copy_input(::typeof(svd_compact), A) = copy_input(svd_full, A)
@@ -362,15 +362,16 @@ function svd_trunc!(A::AbstractMatrix, USVᴴ, alg::TruncatedAlgorithm{<:GPU_Ran
     U, S, Vᴴ = USVᴴ
     _gpu_Xgesvdr!(A, S.diag, U, Vᴴ; alg.alg.kwargs...)
 
-    do_gauge_fix = get(alg.alg.kwargs, :gaugefix, default_gaugefix())::Bool
-    do_gauge_fix && gaugefix!(svd_trunc!, U, Vᴴ)
-
     # TODO: make sure that truncation is based on maxrank, otherwise this might be wrong
-    USVᴴtrunc, ind = truncate(svd_trunc!, (U, S, Vᴴ), alg.trunc)
-    Strunc = diagview(USVᴴtrunc[2])
+    (Utr, Str, Vᴴtr), _ = truncate(svd_trunc!, (U, S, Vᴴ), alg.trunc)
+
     # normal `truncation_error!` does not work here since `S` is not the full singular value spectrum
-    ϵ = sqrt(norm(A)^2 - norm(Strunc)^2) # is there a more accurate way to do this?
-    return USVᴴtrunc..., ϵ
+    ϵ = sqrt(norm(A)^2 - norm(diagview(Str))^2) # is there a more accurate way to do this?
+
+    do_gauge_fix = get(alg.alg.kwargs, :gaugefix, default_gaugefix())::Bool
+    do_gauge_fix && gaugefix!(svd_trunc!, Utr, Vᴴtr)
+
+    return Utr, Str, Vᴴtr, ϵ
 end
 
 function svd_compact!(A::AbstractMatrix, USVᴴ, alg::GPU_SVDAlgorithm)