Probability vectors for https://github.com/tpapp/TransformVariables.jl

ProbVector.jl

export ProbVector
using ArgCheck: @argcheck
using Parameters

import TransformVariables: RealVector,@calltrans, LogJacFlag, 
    VectorTransform, dimension, inverse!, inverse_eltype
using TransformVariables
"""
    (y, r, ℓ) = SIGNATURES
Given ``x ∈ ℝ`` and ``0 ≤ r ≤ 1``, return `(y, r′)` such that
1. ``y + r′ = r``,
2. ``y: 0 ≤ y ≤ r`` is mapped with a bijection from `x`.
`ℓ` is the log Jacobian (whether it is evaluated depends on `flag`).
"""
@inline function l1_remainder_transform(flag::LogJacFlag, x, r)
    z = logistic(x)
    y = z*r
    (y, r-y,
     flag isa NoLogJac ? flag : log(r) - log(z*(1-z)))
end

"""
    (x, r′) = SIGNATURES
Inverse of [`l1_remainder_transform`](@ref) in `x` and `y`.
"""
@inline l1_remainder_inverse(y, r) = logit(y/r), r-y


"""
    ProbVector(n)
Transform `n-1` real numbers to a probability vector of length `n`
"""
@calltrans struct ProbVector <: VectorTransform
    n::Int
    function ProbVector(n::Int)
        @argcheck n ≥ 1 "Dimension should be positive."
        new(n)
    end
end

dimension(t::ProbVector) = t.n - 1

function transform_with(flag::LogJacFlag, t::ProbVector, x::RealVector)
    @unpack n = t
    @argcheck n - 1 == length(x)
    T = extended_eltype(x)
    r = one(T)
    y = Vector{T}(undef, n)
    ℓ = logjac_zero(flag, T)
    index = firstindex(x)
    @inbounds for i in 1:(n - 1)
        xi = x[index]
        index += 1
        y[i], r, ℓi = l1_remainder_transform(flag, xi, r)
        ℓ += ℓi
    end
    y[end] = r
    y, ℓ
end

inverse_eltype(t::ProbVector, y::RealVector) = extended_eltype(y)

function inverse!(x::RealVector, t::ProbVector, y::RealVector)
    @unpack n = t
    @argcheck length(y) == n
    @argcheck length(x) == n - 1
    r = one(eltype(y))
    xi = firstindex(x)
    @inbounds for yi in axes(y, 1)[1:(end-1)]
        x[xi], r = l1_remainder_inverse(y[yi], r)
        xi += 1
    end
    x
end