diff --git a/src/to_vec.jl b/src/to_vec.jl index e5057aa..c42a102 100644 --- a/src/to_vec.jl +++ b/src/to_vec.jl @@ -140,6 +140,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) diff --git a/test/to_vec.jl b/test/to_vec.jl index 2fa1fdb..663ee57 100644 --- a/test/to_vec.jl +++ b/test/to_vec.jl @@ -49,6 +49,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 @@ -116,6 +117,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