Vector-version for PDBijector

Project.toml

@@ -1,6 +1,6 @@
 name = "Bijectors"
 uuid = "76274a88-744f-5084-9051-94815aaf08c4"
-version = "0.12.8"
+version = "0.13.0"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

src/bijectors/corr.jl

@@ -232,6 +232,17 @@ function logabsdetjac(::Inverse{VecCorrBijector}, y::AbstractVector{<:Real})
     return _logabsdetjac_inv_corr(y)
 end
 
+function output_size(::VecCorrBijector, sz::NTuple{2})
+    @assert sz[1] == sz[2]
+    n = sz[1]
+    return (n * (n - 1)) ÷ 2
+end
+
+function output_size(::Inverse{VecCorrBijector}, sz::NTuple{1})
+    n = _triu1_dim_from_length(first(sz))
+    return (n, n)
+end
+
 """
     VecCholeskyBijector <: Bijector
 
@@ -317,6 +328,11 @@ function logabsdetjac(::Inverse{VecCholeskyBijector}, y::AbstractVector{<:Real})
     return _logabsdetjac_inv_chol(y)
 end
 
+output_size(::VecCholeskyBijector, sz::NTuple{2}) = output_size(VecCorrBijector(), sz)
+function output_size(::Inverse{<:VecCholeskyBijector}, sz::NTuple{1})
+    return output_size(inverse(VecCorrBijector()), sz)
+end
+
 """
     function _link_chol_lkj(w)
 

src/bijectors/pd.jl

@@ -40,3 +40,49 @@ end
 function with_logabsdet_jacobian(b::PDBijector, X)
     return transform(b, X), logabsdetjac(b, X)
 end
+
+
+struct PDVecBijector <: Bijector end
+
+function _triu_dim_from_length(d)
+    # (n^2 + n) / 2 = d
+    #   n² + n - 2d = 0
+    #             n = (-1 + sqrt(1 + 8d)) / 2
+    return (-1 + isqrt(1 + 8 * d)) ÷ 2
+end
+
+"""
+    vec_to_triu(x::AbstractVector{<:Real})
+
+Constructs a matrix from a vector `x` by filling the upper triangle.
+"""
+function vec_to_triu(x::AbstractVector)
+    n = _triu_dim_from_length(length(x))
+    X = update_triu_from_vec(x, 0, n)
+    return upper_triangular(X)
+end
+
+transform(::PDVecBijector, X::AbstractMatrix{<:Real}) = pd_vec_link(X)
+pd_vec_link(X) = triu_to_vec(transpose(pd_link(X)), 0)
+
+function transform(::Inverse{PDVecBijector}, y::AbstractVector{<:Real})
+    Y = permutedims(vec_to_triu(y))
+    return transform(inverse(PDBijector()), Y)
+end
+
+logabsdetjac(::PDVecBijector, X::AbstractMatrix{<:Real}) = logabsdetjac(PDBijector(), X)
+
+function with_logabsdet_jacobian(b::PDVecBijector, X)
+    return transform(b, X), logabsdetjac(b, X)
+end
+
+function output_size(::PDVecBijector, sz::Tuple{Int,Int})
+    n = first(sz)
+    d = (n^2 + n) ÷ 2
+    return (d,)
+end
+
+function output_size(::Inverse{PDVecBijector}, sz::Tuple{Int})
+    n = _triu_dim_from_length(first(sz))
+    return (n, n)
+end

src/bijectors/reshape.jl

@@ -25,3 +25,5 @@ end
 inverse(b::Reshape) = Reshape(b.out_shape, b.in_shape)
 
 with_logabsdet_jacobian(b::Reshape, x) = reshape(x, b.out_shape), zero(eltype(x))
+
+output_size(b::Reshape, in_size) = b.out_shape

src/interface.jl

@@ -35,6 +35,26 @@ function logabsdetjac(f::Columnwise, x::AbstractMatrix)
 end
 with_logabsdet_jacobian(f::Columnwise, x::AbstractMatrix) = (f(x), logabsdetjac(f, x))
 
+"""
+    output_size(f, sz)
+
+Returns the output size of `f` given the input size `sz`.
+"""
+output_size(f, sz) = sz
+
+"""
+    output_length(f, len::Int)
+    output_length(f, sz::Tuple)
+
+Returns the output length of `f` given the input length `len` or size `sz`.
+"""
+output_length(f, len::Int) = len
+function output_length(f, len::Tuple)
+    sz = output_size(f, len)
+    @assert length(sz) == 1
+    return first(sz)
+end
+
 ######################
 # Bijector interface #
 ######################

src/transformed_distribution.jl

@@ -1,20 +1,19 @@
+function variateform(d::Distribution, b)
+    sz_in = size(d)
+    sz_out = output_size(b, sz_in)
+    return ArrayLikeVariate{length(sz_out)}
+end
+
+variateform(::MultivariateDistribution, ::Inverse{VecCholeskyBijector}) = CholeskyVariate
+
 # Transformed distributions
 struct TransformedDistribution{D,B,V} <:
-       Distribution{V,Continuous} where {D<:Distribution{V,Continuous},B}
+       Distribution{V,Continuous} where {D<:ContinuousDistribution,B}
     dist::D
     transform::B
 
-    function TransformedDistribution(d::UnivariateDistribution, b)
-        return new{typeof(d),typeof(b),Univariate}(d, b)
-    end
-    function TransformedDistribution(d::MultivariateDistribution, b)
-        return new{typeof(d),typeof(b),Multivariate}(d, b)
-    end
-    function TransformedDistribution(d::MatrixDistribution, b)
-        return new{typeof(d),typeof(b),Matrixvariate}(d, b)
-    end
-    function TransformedDistribution(d::Distribution{CholeskyVariate}, b)
-        return new{typeof(d),typeof(b),CholeskyVariate}(d, b)
+    function TransformedDistribution(d::ContinuousDistribution, b)
+        return new{typeof(d),typeof(b),variateform(d, b)}(d, b)
     end
 end
 
@@ -83,8 +82,8 @@ bijector(d::BoundedDistribution) = bijector_bounded(d)
 const LowerboundedDistribution = Union{Pareto,Levy}
 bijector(d::LowerboundedDistribution) = bijector_lowerbounded(d)
 
-bijector(d::PDMatDistribution) = PDBijector()
-bijector(d::MatrixBeta) = PDBijector()
+bijector(d::PDMatDistribution) = PDVecBijector()
+bijector(d::MatrixBeta) = PDVecBijector()
 
 bijector(d::LKJ) = VecCorrBijector()
 bijector(d::LKJCholesky) = VecCholeskyBijector(d.uplo)
@@ -101,8 +100,8 @@ end
 ##############################
 
 # size
-Base.length(td::Transformed) = length(td.dist)
-Base.size(td::Transformed) = size(td.dist)
+Base.length(td::Transformed) = output_length(td.transform, size(td.dist))
+Base.size(td::Transformed) = output_size(td.transform, size(td.dist))
 
 function logpdf(td::UnivariateTransformed, y::Real)
     x, logjac = with_logabsdet_jacobian(inverse(td.transform), y)

test/bijectors/corr.jl

@@ -33,6 +33,15 @@ using Bijectors: VecCorrBijector, VecCholeskyBijector, CorrBijector
         test_bijector(bvec, x; test_not_identity=d != 1, changes_of_variables_test=false)
 
         test_ad(x -> sum(bvec(bvecinv(x))), yvec)
+
+        # Check that output sizes are computed correctly.
+        tdist = transformed(dist)
+        @test length(tdist) == length(yvec)
+        @test tdist isa MultivariateDistribution
+
+        dist_unconstrained = transformed(MvNormal(zeros(length(tdist)), I), inverse(bvec))
+        @test size(dist_unconstrained) == size(x)
+        @test dist_unconstrained isa MatrixDistribution
     end
 end
 
@@ -60,6 +69,15 @@ end
             # test_bijector is commented out for now, 
             # as isapprox is not defined for ::Cholesky types (the domain of LKJCholesky)
             # test_bijector(b, x; test_not_identity=d != 1, changes_of_variables_test=false)
+
+            # Check that output sizes are computed correctly.
+            tdist = transformed(dist)
+            @test length(tdist) == length(y)
+            @test tdist isa MultivariateDistribution
+
+            dist_unconstrained = transformed(MvNormal(zeros(length(tdist)), I), inverse(b))
+            @test size(dist_unconstrained) == size(x)
+            @test dist_unconstrained isa Distribution{CholeskyVariate,Continuous}
         end
     end
 end

test/bijectors/pd.jl

@@ -1,13 +1,37 @@
 using Bijectors, DistributionsAD, LinearAlgebra, Test
-using Bijectors: PDBijector
+using Bijectors: PDBijector, PDVecBijector
 
 @testset "PDBijector" begin
-    d = 5
-    b = PDBijector()
-    dist = Wishart(d, Matrix{Float64}(I, d, d))
-    x = rand(dist)
-    # NOTE: `PDBijector` technically isn't bijective, and so the default `getjacobian`
-    # used in the ChangesOfVariables.jl tests will fail as the jacobian will have determinant 0.
-    # Hence, we disable those tests.
-    test_bijector(b, x; test_not_identity=true, changes_of_variables_test=false)
+    for d in [2, 5]
+        b = PDBijector()
+        dist = Wishart(d, Matrix{Float64}(I, d, d))
+        x = rand(dist)
+        # NOTE: `PDBijector` technically isn't bijective, and so the default `getjacobian`
+        # used in the ChangesOfVariables.jl tests will fail as the jacobian will have determinant 0.
+        # Hence, we disable those tests.
+        test_bijector(b, x; test_not_identity=true, changes_of_variables_test=false)
+    end
+end
+
+@testset "PDVecBijector" begin
+    for d in [2, 5]
+        b = PDVecBijector()
+        dist = Wishart(d, Matrix{Float64}(I, d, d))
+        x = rand(dist)
+        y = b(x)
+
+        # NOTE: `PDBijector` technically isn't bijective, and so the default `getjacobian`
+        # used in the ChangesOfVariables.jl tests will fail as the jacobian will have determinant 0.
+        # Hence, we disable those tests.
+        test_bijector(b, x; test_not_identity=true, changes_of_variables_test=false)
+
+        # Check that output sizes are computed correctly.
+        tdist = transformed(dist, b)
+        @test length(tdist) == length(y)
+        @test tdist isa MultivariateDistribution
+
+        dist_unconstrained = transformed(MvNormal(zeros(length(tdist)), I), inverse(b))
+        @test size(dist_unconstrained) == size(x)
+        @test dist_unconstrained isa MatrixDistribution
+    end
 end