Vectorize structured matrices densely

Project.toml

@@ -1,6 +1,6 @@
 name = "FiniteDifferences"
 uuid = "26cc04aa-876d-5657-8c51-4c34ba976000"
-version = "0.12.35"
+version = "0.12.36"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/to_vec.jl

@@ -105,12 +105,19 @@ end
 
 function to_vec(x::T) where {T<:LinearAlgebra.HermOrSym}
     x_vec, back = to_vec(Matrix(x))
+    # Structured matrices are vectorized densely, and their `from_vec` projects onto their
+    # structure (e.g. `Diagonal`): reuse it for the parent
+    parent_vec, parent_back = to_vec(parent(x))
+    if length(parent_vec) != length(x_vec)
+        throw(ArgumentError("to_vec($T) requires the parent to be vectorized densely"))
+    end
     function HermOrSym_from_vec(x_vec)
         # Project onto symmetric/hermitian matrices instead of reading only one triangle:
         # that would double the off-diagonals of cotangents, see issue #164
         X = back(x_vec)
         X_sym = (X .+ (T <: Symmetric ? transpose(X) : adjoint(X))) ./ 2
-        return T(X_sym, x.uplo)
+        X_parent = parent_back(first(to_vec(X_sym)))
+        return (T <: Symmetric ? Symmetric : Hermitian)(X_parent, Symbol(x.uplo))
     end
     return x_vec, HermOrSym_from_vec
 end
@@ -124,16 +131,30 @@ function to_vec(x::Diagonal)
 end
 
 function to_vec(x::Tridiagonal)
-    x_vec, back = to_vec((x.dl, x.d, x.du))
-    # Other field (du2) of a Tridiagonal is not part of its value and is really a kind of cache
+    x_vec, back = to_vec(Matrix(x))
     function Tridiagonal_from_vec(x_vec)
-        return Tridiagonal(back(x_vec)...)
+        return Tridiagonal(back(x_vec))
     end
     return x_vec, Tridiagonal_from_vec
 end
 
+function to_vec(x::T) where {T<:SymTridiagonal}
+    x_vec, back = to_vec(Matrix(x))
+    function SymTridiagonal_from_vec(x_vec)
+        # Average the off-diagonals instead of reading only one of them, see issue #164
+        X = back(x_vec)
+        return T(diag(X), (diag(X, 1) .+ diag(X, -1)) ./ 2)
+    end
+    return x_vec, SymTridiagonal_from_vec
+end
 
-
+function to_vec(x::Bidiagonal)
+    x_vec, back = to_vec(Matrix(x))
+    function Bidiagonal_from_vec(x_vec)
+        return Bidiagonal(back(x_vec), Symbol(x.uplo))
+    end
+    return x_vec, Bidiagonal_from_vec
+end
 
 function to_vec(X::Transpose)
     x_vec, back = to_vec(Matrix(X))
@@ -163,10 +184,12 @@ function to_vec(x::Adjoint{<:Any, <:AbstractVector})
     return x_vec, Adjoint_from_vec
 end
 
-function to_vec(X::T) where {T<:PermutedDimsArray}
-    x_vec, back = to_vec(parent(X))
+function to_vec(X::T) where {iperm, T<:PermutedDimsArray{<:Any, <:Any, <:Any, iperm}}
+    # Vectorize in the order of `X`, not of its parent, so `X` and tangents of other types
+    # (e.g. `Array`) are vectorized consistently
+    x_vec, back = to_vec(collect(X))
     function PermutedDimsArray_from_vec(x_vec)
-        X_parent = back(x_vec)
+        X_parent = permutedims(back(x_vec), iperm)
         return T(X_parent)
     end
     return x_vec, PermutedDimsArray_from_vec

test/grad.jl

@@ -221,21 +221,44 @@ end
 # https://github.com/JuliaDiff/FiniteDifferences.jl/issues/164
 @testset "jvp / j′vp(::$Op{$T})" for Op in (Symmetric, Hermitian), T in (Float64, ComplexF64)
     rng, fdm = MersenneTwister(123456), central_fdm(5, 1)
-    @testset "uplo = $uplo" for uplo in (:U, :L)
-        X = Op(randn(rng, T, 3, 3), uplo)
+    @testset "uplo = $uplo, parent = $P" for uplo in (:U, :L),
+        P in (identity, Diagonal, UpperTriangular, LowerTriangular, Tridiagonal)
+
+        X = Op(P(randn(rng, T, 3, 3)), uplo)
 
         # The cotangent of `Matrix(X)` is the projection of the output cotangent onto
-        # symmetric/hermitian matrices, without doubling the off-diagonals
+        # matrices of the same structure, without doubling the off-diagonals
         Ȳ = randn(rng, T, 3, 3)
         X̄ = j′vp(fdm, Matrix, Ȳ, X)[1]
         @test X̄ isa Op
-        @test X̄ ≈ (Ȳ .+ (Op === Symmetric ? transpose(Ȳ) : adjoint(Ȳ))) ./ 2
+        @test X̄ ≈ Op(P((Ȳ .+ (Op === Symmetric ? transpose(Ȳ) : adjoint(Ȳ))) ./ 2), uplo)
 
-        Ẋ = Op(randn(rng, T, 3, 3), uplo)
+        Ẋ = Op(P(randn(rng, T, 3, 3)), uplo)
         @test jvp(fdm, Matrix, (X, Ẋ)) ≈ Ẋ
     end
 end
 
+# Structured matrices are vectorized densely, so they can be combined with dense tangents
+@testset "jvp / j′vp(::$name{$T})" for T in (Float64, ComplexF64), (name, P) in (
+    # orthogonal projections of dense matrices onto the structured ones
+    "Tridiagonal" => Tridiagonal,
+    "SymTridiagonal" => Y -> SymTridiagonal((Y .+ transpose(Y)) ./ 2),
+    "Bidiagonal(:U)" => Y -> Bidiagonal(Y, :U),
+    "Bidiagonal(:L)" => Y -> Bidiagonal(Y, :L),
+    "PermutedDimsArray" => Y -> PermutedDimsArray(permutedims(Y), (2, 1)),
+)
+    rng, fdm = MersenneTwister(123456), central_fdm(5, 1)
+    X = P(randn(rng, T, 4, 4))
+
+    Ȳ = randn(rng, T, 4, 4)
+    X̄ = j′vp(fdm, identity, Ȳ, X)[1]
+    @test X̄ isa typeof(X)
+    @test X̄ ≈ P(Ȳ)
+
+    Ẋ = P(randn(rng, T, 4, 4))
+    @test jvp(fdm, identity, (X, Ẋ)) ≈ Ẋ
+end
+
 using LinearAlgebra
 
 function partial_nan_returning(x)

test/to_vec.jl

@@ -20,6 +20,15 @@ end
 Base.size(x::FillVector) = (x.len,)
 Base.getindex(x::FillVector, n::Int) = x.x
 
+# A dummy square FillMatrix, which is not vectorized densely
+struct FillMatrix <: AbstractMatrix{Float64}
+    x::Float64
+    n::Int
+end
+
+Base.size(x::FillMatrix) = (x.n, x.n)
+Base.getindex(x::FillMatrix, i::Int, j::Int) = x.x
+
 # For testing Tangent{ThreeFields}
 struct ThreeFields
     a
@@ -89,6 +98,9 @@ end
         test_to_vec(UpperTriangular(randn(T, 13, 13)))
         test_to_vec(Diagonal(randn(T, 7)))
         test_to_vec(Tridiagonal(randn(T, 3), randn(T, 4), randn(T, 3)))
+        test_to_vec(SymTridiagonal(randn(T, 4), randn(T, 3)))
+        test_to_vec(Bidiagonal(randn(T, 4), randn(T, 3), :U))
+        test_to_vec(Bidiagonal(randn(T, 4), randn(T, 3), :L))
 
         test_to_vec(DummyType(randn(T, 2, 9)))
         test_to_vec(SVector{2, T}(1.0, 2.0); check_inferred=false)
@@ -105,6 +117,11 @@ end
                 x_vec, back = to_vec(A)
                 @test back(x_vec).uplo == A.uplo
             end
+            @testset "parent = $P" for P in (Diagonal, UpperTriangular, LowerTriangular, Tridiagonal)
+                test_to_vec(Op(P(randn(T, 4, 4)), :U))
+                test_to_vec(Op(P(randn(T, 4, 4)), :L))
+            end
+            @test_throws ArgumentError to_vec(Op(FillMatrix(1.0, 4)))
         end
 
         @testset "$Op" for Op in (Adjoint, Transpose)