add transformation to simplex (#56)

add transform for UnitSimplex

Author: scheidan

src/special_arrays.jl

@@ -1,4 +1,4 @@
-export UnitVector, CorrCholeskyFactor
+export UnitVector, UnitSimplex, CorrCholeskyFactor
 
 ####
 #### building blocks
@@ -77,6 +77,67 @@ function inverse_at!(x::AbstractVector, index, t::UnitVector, y::AbstractVector)
     index
 end
 
+
+####
+#### UnitSimplex
+####
+
+"""
+    UnitSimplex(n)
+
+Transform `n-1` real numbers to a vector of length `n` whose elements are non-negative and sum to one.
+"""
+@calltrans struct UnitSimplex <: VectorTransform
+    n::Int
+    function UnitSimplex(n::Int)
+        @argcheck n ≥ 1 "Dimension should be positive."
+        new(n)
+    end
+end
+
+dimension(t::UnitSimplex) = t.n - 1
+
+function transform_with(flag::LogJacFlag, t::UnitSimplex, x::AbstractVector, index)
+    @unpack n = t
+    T = extended_eltype(x)
+
+    ℓ = logjac_zero(flag, T)
+    stick = one(T)
+    y = Vector{T}(undef, n)
+    @inbounds for i in 1:n-1
+        xi = x[index]
+        index += 1
+        z = logistic(xi - log(n-i))
+        y[i] = z * stick
+
+        if !(flag isa NoLogJac)
+            ℓ += log(stick) - logit_logjac(z)
+        end
+
+        stick *= 1 - z
+    end
+
+    y[end] = stick
+
+    y, ℓ, index
+end
+
+inverse_eltype(t::UnitSimplex, y::AbstractVector) = extended_eltype(y)
+
+function inverse_at!(x::AbstractVector, index, t::UnitSimplex, y::AbstractVector)
+    @unpack n = t
+    @argcheck length(y) == n
+
+    stick = one(eltype(y))
+    @inbounds for i in axes(y, 1)[1:end-1]
+        z = y[i]/stick
+        x[index] = logit(z) + log(n-i)
+        stick -= y[i]
+        index += 1
+    end
+    index
+end
+
 ####
 #### correlation cholesky factor
 ####

test/runtests.jl

@@ -103,6 +103,29 @@ end
     end
 end
 
+@testset "to unit simplex" begin
+    @testset "dimension checks" begin
+        S = UnitSimplex(3)
+        x = zeros(3)               # incorrect
+        @test_throws ArgumentError S(x)
+        @test_throws ArgumentError transform(S, x)
+        @test_throws ArgumentError transform_and_logjac(S, x)
+    end
+
+    @testset "consistency checks" begin
+        for K in 1:10
+            t = UnitSimplex(K)
+            @test dimension(t) == K - 1
+            if K > 1
+                test_transformation(t, y -> (sum(y) ≈ 1) & (all(y.>=0)),
+                                    vec_y = y -> y[1:(end-1)])
+            end
+            x = zeros(dimension(t))
+            @test transform(t, x) ≈ 1 ./ fill(K, K)
+        end
+    end
+end
+
 @testset "to correlation cholesky factor" begin
     @testset "dimension checks" begin
         C = CorrCholeskyFactor(3)