Implement to_vec for matrix factorizations

src/to_vec.jl

@@ -162,11 +162,24 @@ function to_vec(x::Cholesky)
  return x_vec, Cholesky_from_vec
 end
 
-function to_vec(x::S) where {S <: Union{LinearAlgebra.QRCompactWYQ, LinearAlgebra.QRCompactWY}}
- x_vec, back = to_vec([x.factors, x.T])
+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, T = back(v)
- return S(factors, T)
+ factors, Tback = back(v)
+ return S(factors, Tback)
  end
  return x_vec, QRCompact_from_vec
 end

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
@@ -117,12 +118,27 @@ end
  end
 
  @testset "Factorizations" begin
- for dims in [(5, 5), (4, 6), (7, 3)]
+ # (100, 100) is needed to test for the NaNs that can appear in the
+ # qr(M).T matrix
+ for dims in [(5, 5), (4, 6), (7, 3), (100, 100)]
  M = randn(T, dims...)
  P = M * M' + I # Positive definite matrix
- test_to_vec(svd(M); check_inferred = true)
- test_to_vec(qr(M))
+ 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.Q == F.Q
+ @test F.R == F.R
+ @test F_vec == first(to_vec(F))
  end
  end