Implement to_vec for matrix factorizations (#128)

* Implement to_vec for factorizations

* Fix inferrence for svd

* Fix problem from qr(X).T undef values

* Fix tests

* Reduce number of tests

Co-authored-by: Simeon Schaub <[email protected]>

Co-authored-by: Simeon Schaub <[email protected]>

Author: Kolaru

src/to_vec.jl

@@ -156,6 +156,51 @@ function to_vec(X::T) where {T<:PermutedDimsArray}
     return x_vec, PermutedDimsArray_from_vec
 end
 
+# Factorizations
+
+function to_vec(x::F) where {F <: SVD}
+    # Convert the vector S to a matrix so we can work with a vector of matrices
+    # only and inferrence work
+    v = [x.U, reshape(x.S, length(x.S), 1), x.Vt]
+    x_vec, back = to_vec(v)
+    function SVD_from_vec(v)
+        U, Smat, Vt = back(v)
+        return F(U, vec(Smat), Vt)
+    end
+    return x_vec, SVD_from_vec
+end
+
+function to_vec(x::Cholesky)
+    x_vec, back = to_vec(x.factors)
+    function Cholesky_from_vec(v)
+        return Cholesky(back(v), x.uplo, x.info)
+    end
+    return x_vec, Cholesky_from_vec
+end
+
+function to_vec(x::S) where {U, S <: Union{LinearAlgebra.QRCompactWYQ{U}, LinearAlgebra.QRCompactWY{U}}}
+    # x.T is composed of upper triangular blocks. The subdiagonals elements
+    # of the blocks are abitrary. We make sure to set all of them to zero
+    # to avoid NaN.
+    blocksize, cols = size(x.T)
+    T = zeros(U, blocksize, cols)
+
+    for i in 0:div(cols - 1, blocksize)
+        used_cols = i * blocksize
+        n = min(blocksize, cols - used_cols)
+        T[1:n, (1:n) .+ used_cols] = UpperTriangular(view(x.T, 1:n, (1:n) .+ used_cols))
+    end
+
+    x_vec, back = to_vec([x.factors, T])
+
+    function QRCompact_from_vec(v)
+        factors, Tback = back(v)
+        return S(factors, Tback)
+    end
+
+    return x_vec, QRCompact_from_vec
+end
+
 # Non-array data structures
 
 function to_vec(x::Tuple)

test/to_vec.jl

@@ -61,6 +61,7 @@ function test_to_vec(x::T; check_inferred=true) where {T}
     x_vec, back = to_vec(x)
     @test x_vec isa Vector
     @test all(s -> s isa Real, x_vec)
+    @test all(!isnan, x_vec)
     check_inferred && @inferred back(x_vec)
     @test x == back(x_vec)
     return nothing
@@ -128,6 +129,47 @@ end
             )
         end
 
+        @testset "Factorizations" begin
+            # (100, 100) is needed to test for the NaNs that can appear in the
+            # qr(M).T matrix
+            for dims in [(7, 3), (100, 100)]
+                M = randn(T, dims...)
+                P = M * M' + I  # Positive definite matrix
+                test_to_vec(svd(M))
+                test_to_vec(cholesky(P))
+
+                # Special treatment for QR since it is represented by a matrix
+                # with some arbirtrary values.
+                F = qr(M)
+                @inferred to_vec(F)
+                F_vec, back = to_vec(F)
+                @test F_vec isa Vector
+                @test all(s -> s isa Real, F_vec)
+                @test all(!isnan, F_vec)
+                @inferred back(F_vec)
+                F_back = back(F_vec)
+                @test F_back.Q == F.Q
+                @test F_back.R == F.R
+
+                # Make sure the result is consistent despite the arbitrary
+                # values in F.T.
+                @test first(to_vec(F)) == first(to_vec(F))
+
+                # Test F.Q as well since it has a special type. Since it is
+                # represented by the same T and factors matrices than F
+                # it needs the same special treatment.
+                Q = F.Q
+                @inferred to_vec(Q)
+                Q_vec, back = to_vec(Q)
+                @test Q_vec isa Vector
+                @test all(s -> s isa Real, Q_vec)
+                @test all(!isnan, Q_vec)
+                @inferred back(Q_vec)
+                Q_back = back(Q_vec)
+                @test Q_back == Q
+            end
+        end
+
         @testset "Tuples" begin
             test_to_vec((5, 4))
             test_to_vec((5, randn(T, 5)); check_inferred = VERSION ≥ v"1.2") # broken on Julia 1.6.0, fixed on 1.6.1