Constrained structs with automatic parameter handling

One macro, two backends: TransformVariables.jl and Bijectors.jl
Abstract

A struct whose fields carry constraints (a positive scalar, a vector on the simplex, a correlation matrix) needs a constructor that rejects invalid values, and a bijection between the constrained set and $\mathbb{R}^p$ so that optimizers and samplers can work in unconstrained coordinates. This notebook generates both from the struct definition, over TransformVariables.jl and over Bijectors.jl, and checks each generated transform against the definition of the constraint it came from.

The question

From a Julia Discourse thread, the wanted syntax is roughly this:

@probably_a_macro struct MyStruct{T}
    a::T          # a must be positive
    b::Vector{T}  # b must belong to the simplex
    C::Matrix{T}  # must be a correlation matrix
end

with a strict constructor, an intrinsic_dimension, a map $\mathbb{R}^p \to$ MyStruct with its Jacobian, an unconstrain going the other way, and nesting.

This notebook builds @constrained, which gives all of that:

@constrained struct MixtureParams{N,P,T}
    σ::T         ~ positive
    w::Vector{T} ~ simplex(N)
    C::Matrix{T} ~ corrmatrix(P)
end

@constrained struct Model{N,P,T}       # nesting composes
    marginal::MixtureParams{N,P,T}
    ν::T ~ negative
end

intrinsic_dimension(Model{3,4})        # 10
m = constrain(Model{3,4}, randn(10))   # valid by construction
unconstrain(m)                         # back to R^10
constrain_and_logjac(Model{3,4}, x)

Most of the pieces exist. TransformVariables.jl and Bijectors.jl both have the bijections and the log-Jacobians, and TransformVariables can already produce a struct rather than a NamedTuple. What is missing is the declaration. The macro below is written over both packages, to compare them.

using LinearAlgebra, Random, Statistics
import TransformVariables as TV
import Bijectors as BJ
import ConstructionBase
import ForwardDiff, Optim, LineSearches
using Distributions: MvNormal, Exponential, Uniform, Dirichlet, LKJ, LKJCholesky

(; TransformVariables = pkgversion(TV), Bijectors = pkgversion(BJ))
(TransformVariables = v"0.8.26", Bijectors = v"0.16.3")

Describing the constraints

The macro should not know about either package. It reads the struct body and emits one descriptor value per field, and a backend turns descriptors into transforms. A third backend is then new methods rather than a change to the macro.

Descriptors also carry the validation predicate that the strict constructor calls.

abstract type Constraint end

struct Unbounded <: Constraint end
struct Positive  <: Constraint end
struct Negative  <: Constraint end
struct Interval  <: Constraint; lo::Float64; hi::Float64; end
struct Simplex   <: Constraint; n::Int; end
"The Cholesky factor `U` of a correlation matrix: upper triangular, `U'U` has unit diagonal."
struct CorrFactor <: Constraint; n::Int; end
"The correlation matrix itself, carried as a dense symmetric matrix."
struct CorrMatrix <: Constraint; n::Int; end
"A field whose type is itself a constrained struct."
struct Nested <: Constraint; T::Type; end

satisfies(::Unbounded, v) = true
satisfies(::Positive, v)  = v > zero(v)
satisfies(::Negative, v)  = v < zero(v)
satisfies(c::Interval, v) = c.lo < v < c.hi
satisfies(c::Simplex, v)  = length(v) == c.n && all((0), v) && isapprox(sum(v), 1; atol = 1e-8)
satisfies(c::CorrFactor, v) = size(v) == (c.n, c.n) && istriu(v) &&
                              all(j -> isapprox(sum(abs2, @view v[:, j]), 1; atol = 1e-8), 1:c.n)
satisfies(c::CorrMatrix, v) = size(v) == (c.n, c.n) && issymmetric(v) &&
                              all(x -> isapprox(x, 1; atol = 1e-8), diag(v)) && isposdef(v)
satisfies(c::Nested, v) = v isa c.T

describe(::Unbounded) = "unconstrained"
describe(::Positive)  = "positive"
describe(::Negative)  = "negative"
describe(c::Interval) = "in ($(c.lo), $(c.hi))"
describe(c::Simplex)  = "on the $(c.n)-element unit simplex"
describe(c::CorrFactor) = "a $(c.n)×$(c.n) correlation Cholesky factor"
describe(c::CorrMatrix) = "a $(c.n)×$(c.n) correlation matrix"
describe(c::Nested)     = "a valid $(nameof(c.T))"

nothing

What a constraint is

Each descriptor names three things:

  1. a subset $\mathcal{S}$ of a vector space, tested by satisfies
  2. its dimension $d$ as a smooth manifold
  3. a diffeomorphism $f : \mathbb{R}^d \to \mathcal{S}$, together with $\log|\det Df|$

simplex(n) is $\{x \in \mathbb{R}^n : x_i > 0,\ \sum_i x_i = 1\}$ with $d = n - 1$. corrmatrix(n) is $\{C \in \mathbb{R}^{n \times n} : C = C^\top,\ C \succ 0,\ C_{ii} = 1\}$ with $d = n(n-1)/2$. corrfactor(n) is the set of upper triangular $U$ with positive diagonal and unit-norm columns, which is the image of corrmatrix(n) under the Cholesky factorisation and has the same dimension.

Both maps are onto the relative interior. A simplex vector produced by constrain has every $x_i > 0$, and no $x$ maps to a point on the boundary.

The three properties are checked numerically further down.

What TransformVariables already gives you

as(::Type{S}, inner) wraps a transform so it produces an S instead of a NamedTuple. That is most of the struct binding, with two requirements.

It builds the value through ConstructionBase.constructorof(S), whose default strips the type parameters, so a struct carrying its sizes in type parameters has to say so. It also constructs by keyword. A macro can write both lines.

Nesting then needs no extra machinery. A nested field's entry in the outer transform is the inner struct's own transform, and the dimensions add up.

"""
    constraints_of(::Type{S}) -> NamedTuple of Constraint, one per field

The only method the macro generates that a backend reads.
"""
function constraints_of end

"""
Run `inner`, then apply `fwd` to its result. `bwd` undoes that, `bwdtype` is the same map on
types (which `TransformVariables` needs to pick the inverse element type), and `Δlogjac` is
the log-Jacobian of `fwd` itself, evaluated at the inner value.
"""
struct Derived{Tr,F,G,H,J} <: TV.VectorTransform
    inner::Tr
    fwd::F
    bwd::G
    bwdtype::H
    Δlogjac::J
end

TV.dimension(t::Derived) = TV.dimension(t.inner)
_addjac(::TV.NoLogJac, ℓ, t, y) =_addjac(::TV.LogJac, ℓ, t, y)   =+ t.Δlogjac(y)
function TV.transform_with(flag::TV.LogJacFlag, t::Derived, x, index)
    y, ℓ, index′ = TV.transform_with(flag, t.inner, x, index)
    t.fwd(y), _addjac(flag, ℓ, t, y), index′
end
TV.inverse_eltype(t::Derived, ::Type{T}) where {T} = TV.inverse_eltype(t.inner, t.bwdtype(T))
TV.inverse_at!(x::AbstractVector, index, t::Derived, y) =
    TV.inverse_at!(x, index, t.inner, t.bwd(y))

# Constraint -> TransformVariables transform.
tv(::Unbounded)   = TV.asℝ
tv(::Positive)    = TV.asℝ₊
tv(::Negative)    = TV.asℝ₋
tv(c::Interval)   = TV.as(Real, c.lo, c.hi)
tv(c::Simplex)    = TV.UnitSimplex(c.n)
tv(c::CorrFactor) = TV.corr_cholesky_factor(c.n)

# corr_cholesky_factor gives the factor U and the log-Jacobian of x ↦ U. Going on to C = U'U
# is a second change of variables, whose log-Jacobian is derived and checked below.
corr_factor_to_matrix_logjac(U, n) = sum(i -> (n - i) * log(U[i, i]), 2:(n - 1);
                                         init = zero(eltype(U)))

tv(c::CorrMatrix) = Derived(TV.corr_cholesky_factor(c.n),
                            U -> Matrix(Symmetric(U'U)),
                            C -> UpperTriangular(cholesky(Symmetric(C)).U),
                            T -> UpperTriangular{eltype(T),Matrix{eltype(T)}},
                            U -> corr_factor_to_matrix_logjac(U, c.n))
tv(c::Nested)     = tv_transform(c.T)

"The whole struct type as one transform from ℝᵖ."
tv_transform(::Type{S}) where {S} = TV.as(_key(S), TV.as(map(tv, constraints_of(S))))

intrinsic_dimension(::Type{S}) where {S} = TV.dimension(tv_transform(S))
constrain(::Type{S}, x::AbstractVector) where {S} = TV.transform(tv_transform(S), x)
constrain_and_logjac(::Type{S}, x::AbstractVector) where {S} =
    TV.transform_and_logjac(tv_transform(S), x)
unconstrain(s::S) where {S} = TV.inverse(tv_transform(S), s)

nothing

The macro

The macro rests on two conventions.

Sizes live in type parameters, so intrinsic_dimension is a function of the type alone. That is what makes nesting compose: a nested field contributes the dimension of its own type, with nothing to thread through.

The element type is the last parameter, so Name{sizes...} names the type with the element left open. A transform can then build the struct at whatever number type it is handed, which is what lets ForwardDiff through.

The constraint goes after the field with ~. :: binds tighter, so the macro receives the field declaration and the constraint as two separate arguments.

const ANNOTATIONS = Dict{Symbol,Symbol}(
    :unbounded  => :Unbounded,
    :positive   => :Positive,
    :negative   => :Negative,
    :interval   => :Interval,
    :simplex    => :Simplex,
    :corrmatrix => :CorrMatrix,
    :corrfactor => :CorrFactor,
)

"Marks the constructor that skips validation: the one a transform calls."
struct Unchecked end

"What `ConstructionBase.constructorof` returns. It builds through the unchecked constructor."
struct Build{S} end
(::Build{S})(; kw...) where {S} = S(Unchecked(); kw...)

"`MixtureParams{3,4,Float64}` → `MixtureParams{3,4}`: the type with the element left open."
function _key end

"A field with no `~`: nested if its type is itself constrained, otherwise free."
_auto_constraint(::Type{T}) where {T} =
    hasmethod(constraints_of, Tuple{Type{T}}) ? Nested(_key(T)) : Unbounded()

function _check(::Type{S}, values::NamedTuple) where {S}
    cs = constraints_of(S)
    for k in keys(cs)
        satisfies(cs[k], values[k]) ||
            throw(ArgumentError("$(nameof(S)): field `$k` must be $(describe(cs[k]))"))
    end
end

_pname(p) = Meta.isexpr(p, :<:) ? p.args[1] : p
_mentions(ex::Symbol, syms) = ex in syms
_mentions(ex::Expr, syms)   = any(a -> _mentions(a, syms), ex.args)
_mentions(::Any, syms)      = false
_drop_last(ex::Expr) = length(ex.args) == 2 ? ex.args[1] : Expr(:curly, ex.args[1], ex.args[2:(end - 1)]...)

_descriptor(a::Symbol, _, _) =
    Expr(:call, get(() -> error("@constrained: unknown constraint `$a`"), ANNOTATIONS, a))
function _descriptor(a::Expr, _, _)
    Meta.isexpr(a, :call) || error("@constrained: cannot read constraint `$a`")
    f = a.args[1]
    Expr(:call, get(() -> error("@constrained: unknown constraint `$f`"), ANNOTATIONS, f),
         a.args[2:end]...)
end
function _descriptor(::Nothing, ftype, elem)
    if Meta.isexpr(ftype, :curly) && ftype.args[end] === elem
        :(_auto_constraint($(_drop_last(ftype))))          # a nested constrained struct
    elseif ftype === elem
        :(Unbounded())                                      # a free scalar
    elseif elem !== nothing && _mentions(ftype, (elem,))
        error("@constrained: field of type `$ftype` needs an explicit constraint")
    else
        :(_auto_constraint($ftype))
    end
end

"""
    @constrained struct Name{Sizes..., T} ... end

Write out the struct with a validating constructor, plus the methods the backends read.

The last type parameter is the element type and the earlier ones are sizes, so that
`Name{sizes...}` names the type with the element left open. That is what lets a transform
build the struct at whatever number type it is handed, including `ForwardDiff.Dual`.
"""
macro constrained(expr)
    Meta.isexpr(expr, :struct) || error("@constrained expects a struct definition")
    ismutable, sig, body = expr.args
    decl = Meta.isexpr(sig, :<:) ? sig.args[1] : sig
    name, params = Meta.isexpr(decl, :curly) ? (decl.args[1], decl.args[2:end]) : (decl, [])
    pnames = _pname.(params)

    sizeparams = isempty(params) ? [] : params[1:(end - 1)]
    sizenames  = _pname.(sizeparams)
    elem       = isempty(pnames) ? nothing : pnames[end]

    Sfull = isempty(pnames) ? name : Expr(:curly, name, pnames...)
    Ssize = isempty(sizenames) ? name : Expr(:curly, name, sizenames...)
    anytype(S) = Expr(:(::), Expr(:curly, :Type, Expr(:<:, S)))
    allwhere(ex)  = isempty(params) ? ex : Expr(:where, ex, params...)
    sizewhere(ex) = isempty(sizeparams) ? ex : Expr(:where, ex, sizeparams...)

    fields, ftypes, descs = Symbol[], Any[], Any[]
    sized = Dict{Symbol,Any}()                     # size parameter => how to read it off a value
    for line in body.args
        line isa LineNumberNode && continue
        annot = nothing
        if Meta.isexpr(line, :call) && line.args[1] === :~
            annot, line = line.args[3], line.args[2]
        end
        Meta.isexpr(line, :(::)) || error("@constrained: `$line` needs a field type")
        fname, ftype = line.args
        push!(fields, fname); push!(ftypes, ftype)
        push!(descs, _descriptor(annot, ftype, elem))
        if Meta.isexpr(annot, :call) && length(annot.args) == 2 &&
           annot.args[2] isa Symbol && annot.args[2] in sizenames
            sized[annot.args[2]] =
                annot.args[1] === :simplex ? :(length($fname)) : :(size($fname, 1))
        end
    end

    namedtuple(vals) = Expr(:tuple, (Expr(:(=), f, v) for (f, v) in zip(fields, vals))...)
    typedargs = [Expr(:(::), f, t) for (f, t) in zip(fields, ftypes)]
    newcall = isempty(pnames) ? Expr(:call, :new, fields...) :
                                Expr(:call, Expr(:curly, :new, pnames...), fields...)

    out = Expr(:block,
        Expr(:struct, ismutable, sig, Expr(:block, typedargs...,
             # unchecked: a transform's output is in the constraint set by construction
             Expr(:(=), allwhere(Expr(:call, Sfull, Expr(:(::), :Unchecked), fields...)), newcall),
             # the strict one
             Expr(:function, allwhere(Expr(:call, Sfull, fields...)),
                  quote
                      _check($Ssize, $(namedtuple(fields)))
                      $newcall
                  end))),
        Expr(:function, sizewhere(Expr(:call, Ssize, Expr(:parameters, fields...))),
             Expr(:call, Ssize, fields...)),
        Expr(:(=), sizewhere(Expr(:call, :(ConstructionBase.constructorof), anytype(Ssize))),
             Expr(:call, Expr(:curly, :Build, Ssize))),
        Expr(:(=), sizewhere(Expr(:call, :_key, anytype(Ssize))), Ssize),
        Expr(:(=), sizewhere(Expr(:call, :constraints_of, anytype(Ssize))), namedtuple(descs)),
        Expr(:function, sizewhere(Expr(:call, Ssize, Expr(:parameters, fields...),
                                       Expr(:(::), :u, :Unchecked))),
             Expr(:call, Ssize, :u, fields...)))

    # Name{sizes...}(values...): take the element type from the values.
    if !isempty(params)
        push!(out.args, Expr(:function, allwhere(Expr(:call, Ssize, typedargs...)),
                             Expr(:call, Sfull, fields...)))
        push!(out.args, Expr(:function,
                             allwhere(Expr(:call, Ssize, Expr(:(::), :u, :Unchecked), typedargs...)),
                             Expr(:call, Sfull, :u, fields...)))
    end

    # Name(values...): also read the sizes off the values, when no field type mentions one.
    if !isempty(sizenames) && length(sized) == length(sizenames) &&
       !any(t -> _mentions(t, sizenames), ftypes)
        callsig = Expr(:where, Expr(:call, name, typedargs...), params[end])
        target = Expr(:curly, name, (sized[p] for p in sizenames)..., elem)
        push!(out.args, Expr(:function, callsig, Expr(:call, target, fields...)))
    end

    esc(out)
end
Main.NB.@constrained

Using it

N and P are the sizes, and T is the element type.

@constrained struct MixtureParams{N,P,T}
    σ::T         ~ positive
    w::Vector{T} ~ simplex(N)
    C::Matrix{T} ~ corrmatrix(P)
end

constraints_of(MixtureParams{3,4})
(σ = Main.NB.Positive(), w = Main.NB.Simplex(3), C = Main.NB.CorrMatrix(4))
S = MixtureParams{3,4}
x = randn(Xoshiro(20260920), intrinsic_dimension(S))
θ = constrain(S, x)
Main.NB.MixtureParams{3, 4, Float64}(0.8105296307717669, [0.4466740480638376, 0.37695378040497357, 0.17637217153118884], [1.0 -0.8613688437821857 0.8869776582238944 -0.06360536049387613; -0.8613688437821857 1.0 -0.5487990783806369 -0.24048035509509635; 0.8869776582238944 -0.5487990783806369 1.0 -0.39252633374234797; -0.06360536049387613 -0.24048035509509635 -0.39252633374234797 1.0])
(; p               = intrinsic_dimension(S),
   weights_sum     = sum.w),
   roundtrip_error = maximum(abs, unconstrain(θ) .- x))
(p = 9, weights_sum = 1.0, roundtrip_error = 1.3322676295501878e-15)

The constructor is the strict one. N and P are filled in from the values, because no field type mentions them.

rejected(args...) = try MixtureParams(args...); "accepted" catch e; sprint(showerror, e) end

ok = MixtureParams(2.0, [0.5, 0.3, 0.2], Matrix(1.0I, 4, 4))

(; inferred      = typeof(ok),
   negative_σ    = rejected(-2.0, [0.5, 0.3, 0.2], Matrix(1.0I, 4, 4)),
   w_off_simplex = rejected(2.0, [0.5, 0.3, 0.3], Matrix(1.0I, 4, 4)),
   not_a_corr    = rejected(2.0, [0.5, 0.3, 0.2], [1.0 2.0; 2.0 1.0]))
(inferred = Main.NB.MixtureParams{3, 4, Float64}, negative_σ = "ArgumentError: MixtureParams: field `σ` must be positive", w_off_simplex = "ArgumentError: MixtureParams: field `w` must be on the 3-element unit simplex", not_a_corr = "ArgumentError: MixtureParams: field `C` must be a 2×2 correlation matrix")

Nesting

An unannotated field whose type is itself constrained becomes a Nested constraint, which resolves to that type's own transform. Neither the macro nor the backends treat nesting as a special case.

Here the field types mention N and P, so the macro does not emit the convenience constructor and the type parameters are written out.

@constrained struct Model{N,P,T}
    marginal::MixtureParams{N,P,T}
    ν::T ~ negative
end

M  = Model{3,4}
xm = randn(Xoshiro(7), intrinsic_dimension(M))
m  = constrain(M, xm)

(; inner           = intrinsic_dimension(MixtureParams{3,4}),
   outer           = intrinsic_dimension(M),
   constraints     = constraints_of(M),
   ν               = m.ν,
   roundtrip_error = maximum(abs, unconstrain(m) .- xm))
(inner = 9, outer = 10, constraints = (marginal = Main.NB.Nested(Main.NB.MixtureParams{3, 4}), ν = Main.NB.Negative()), ν = -0.5602199410015952, roundtrip_error = 6.661338147750939e-16)

What the vector holds

constraints_of is recursive, so the layout of x can be read straight off the type. The sliders change the sizes.

"Flatten a constrained type into the leaf constraints, in the order they occupy `x`."
function leaves(::Type{S}, prefix = "") where {S}
    out = Pair{String,Int}[]
    for (k, c) in pairs(constraints_of(S))
        label = isempty(prefix) ? string(k) : "$prefix.$k"
        c isa Nested ? append!(out, leaves(c.T, label)) : push!(out, label => TV.dimension(tv(c)))
    end
    out
end

leaves(Model{3,4})
4-element Vector{Pair{String, Int64}}:
 "marginal.σ" => 1
 "marginal.w" => 2
 "marginal.C" => 6
          "ν" => 1

The log-Jacobian

constrain_and_logjac returns the log absolute determinant of the map from x to the parameter. It can be checked against a numerical Jacobian, provided you differentiate with respect to coordinates that are actually free: one number for a scalar, the first $N-1$ entries of a simplex vector, the strict upper triangle of a correlation Cholesky factor.

free(::Union{Unbounded,Positive,Negative,Interval}, v) = [v]
free(::Simplex, v)   = v[1:end-1]
free(c::CorrFactor, U) = [U[i, j] for j in 2:c.n for i in 1:(j - 1)]
free(c::CorrMatrix, C) = [C[i, j] for j in 2:c.n for i in 1:(j - 1)]
free(::Nested, s)      = freecoords(s)

"The `p` free numbers inside a constrained value."
freecoords(s::S) where {S} =
    reduce(vcat, [free(c, getfield(s, k)) for (k, c) in pairs(constraints_of(S))])

"log|det J| of x ↦ freecoords(constrain(S, x)), by automatic differentiation."
numeric_logjac(::Type{S}, x) where {S} =
    first(logabsdet(ForwardDiff.jacobian(z -> freecoords(constrain(S, z)), x)))

nothing
@constrained struct FactorParams{N,P,T}
    σ::T                            ~ positive
    w::Vector{T}                    ~ simplex(N)
    U::UpperTriangular{T,Matrix{T}} ~ corrfactor(P)
end

F  = FactorParams{3,4}
xf = randn(Xoshiro(11), intrinsic_dimension(F))
_, ℓ = constrain_and_logjac(F, xf)

(; reported = ℓ, numeric = numeric_logjac(F, xf), difference =- numeric_logjac(F, xf))
(reported = -10.494613610530607, numeric = -10.494613610530605, difference = -1.7763568394002505e-15)

The correlation matrix needs a second term

corr_cholesky_factor produces the factor $U$ and reports the log-Jacobian of $x \mapsto U$. A field declared as the correlation matrix holds $C = U'U$, which is a further change of variables, so it needs the extra term $\log|\det \partial C / \partial U|$.

On the free coordinates that Jacobian is block lower triangular. Ordering both sides by column and writing $C_{ij} = \sum_{k \le i} U_{ki} U_{kj}$ for $i < j$, the block for column $j$ is $\partial(C_{1j},\dots,C_{j-1,j}) / \partial(U_{1j},\dots,U_{j-1,j}) = U[1{:}j-1,\,1{:}j-1]^\top$, whose determinant is $\prod_{i<j} U_{ii}$. Multiplying over $j$,

$$\log\left|\det \frac{\partial C}{\partial U}\right| = \sum_{i=2}^{n-1} (n - i)\,\log U_{ii}.$$

tv(::CorrMatrix) adds that term, so corrmatrix reports the log-Jacobian of $x \mapsto C$. Against automatic differentiation:

"Rebuild the correlation Cholesky factor from its strict upper triangle."
function factor_from_free(u, n)
    U = zeros(eltype(u), n, n)
    U[1, 1] = one(eltype(u))
    k = 0
    for j in 2:n
        for i in 1:(j - 1)
            U[i, j] = u[k += 1]
        end
        U[j, j] = sqrt(1 - sum(abs2, @view U[1:(j - 1), j]))
    end
    UpperTriangular(U)
end

"log|det ∂C/∂U| on the free coordinates, by automatic differentiation."
function ad_corr_gap(U, n)
    J = ForwardDiff.jacobian(free(CorrFactor(n), U)) do u
        V = factor_from_free(u, n)
        free(CorrMatrix(n), Matrix(Symmetric(V'V)))
    end
    first(logabsdet(J))
end

let rng = Xoshiro(3)
    map(2:6) do n
        U = cholesky(Symmetric(constrain(MixtureParams{2,n},
                                         randn(rng, intrinsic_dimension(MixtureParams{2,n}))).C)).U
        ad, closed = ad_corr_gap(U, n), corr_factor_to_matrix_logjac(U, n)
        (; n, ad, closed, difference = ad - closed)
    end
end
5-element Vector{@NamedTuple{n::Int64, ad::Float64, closed::Float64, difference::Float64}}:
 (n = 2, ad = 0.0, closed = 0.0, difference = 0.0)
 (n = 3, ad = -0.9793131474781251, closed = -0.9793131474781251, difference = 0.0)
 (n = 4, ad = -0.4045106258472887, closed = -0.4045106258472889, difference = 2.220446049250313e-16)
 (n = 5, ad = -1.2297639112162129, closed = -1.2297639112162135, difference = 6.661338147750939e-16)
 (n = 6, ad = -12.653509369966464, closed = -12.653509369966446, difference = -1.7763568394002505e-14)
let (_, ℓ) = constrain_and_logjac(S, x), ν = numeric_logjac(S, x)   # S has the C::Matrix field
    (; reported = ℓ, numeric = ν, difference =- ν)
end
(reported = -13.37164379526016, numeric = -13.371643795260162, difference = 1.7763568394002505e-15)

Checking the contract

Three properties, over random draws, for every constraint in the vocabulary.

in_set is satisfies(c, f(u)), so the image really is inside $\mathcal{S}$. left_inverse is $\max \|f^{-1}(f(u)) - u\|_\infty$ and right_inverse is the same the other way round, which together make $f$ a bijection onto its image. logjac compares the reported $\log|\det Df|$ against a Jacobian taken by automatic differentiation in the free coordinates.

"Check the three properties for one constraint over `n` random draws."
function contract(c; n = 400, scale = 1.5, rng = Xoshiro(5))
    t = TV.as((v = tv(c),))
    d = TV.dimension(t)
    in_set, left, right, jac = true, 0.0, 0.0, 0.0
    for _ in 1:n
        u = scale .* randn(rng, d)
        y, ℓ = TV.transform_and_logjac(t, u)
        in_set &= satisfies(c, y.v)
        u′ = TV.inverse(t, y)
        left  = max(left, maximum(abs, u′ .- u))
        right = max(right, maximum(abs, free(c, TV.transform(t, u′).v) .- free(c, y.v)))
        ĵ = first(logabsdet(ForwardDiff.jacobian(z -> free(c, TV.transform(t, z).v), u)))
        jac = max(jac, abs(ℓ - ĵ))
    end
    (; constraint = describe(c), d, in_set,
       left_inverse = left, right_inverse = right, logjac = jac)
end

slate_table([contract(c) for c in (Positive(), Negative(), Interval(-2.0, 3.0),
                                   Simplex(5), CorrFactor(4), CorrMatrix(4))];
            format = (left_inverse  = :scientific,
                      right_inverse = :scientific,
                      logjac        = :scientific),
            align  = (constraint = :left,))
constraintdin_setleft_inverseright_inverselogjac
positive1true1.11e-160e01.11e-16
negative1true1.11e-160e01.11e-16
in (-2.0, 3.0)1true5.77e-158.88e-161.33e-15
on the 5-element unit simplex4true1.09e-131.67e-163.55e-15
a 4×4 correlation Cholesky factor6true3.88e-117.22e-162.42e-13
a 4×4 correlation matrix6true7.04e-113.44e-152.42e-13

The simplex at scale

The forward map lands strictly inside the simplex at every size tried, with the sum correct to $10^{-15}$. The inverse is the direction that degrades. Stick-breaking divides each component by the mass still unallocated, so once a component falls far enough below that remainder the ratio stops being representable and logit cannot recover the coordinate.

scale is the standard deviation of the draws in $\mathbb{R}^d$. At scale = 1 the round trip is exact to $10^{-11}$ out to 200 components. At scale = 3 the smallest component reaches $10^{-35}$ and the inverse stops working, which is a limit of the floating-point type rather than of the map. For a simplex with many components, carry the unconstrained vector instead of recovering it from the constrained value.

"Worst case over random draws of the three things `simplex(n)` claims."
function simplex_report(n; draws = 2000, scale = 3.0, rng = Xoshiro(17))
    t = TV.UnitSimplex(n)
    lo, sumerr, rt, failed = Inf, 0.0, 0.0, 0
    for _ in 1:draws
        u = scale .* randn(rng, TV.dimension(t))
        v = TV.transform(t, u)
        lo     = min(lo, minimum(v))
        sumerr = max(sumerr, abs(sum(v) - 1))
        try
            rt = max(rt, maximum(abs, TV.inverse(t, v) .- u))
        catch
            failed += 1                       # stick-breaking underflow
        end
    end
    (; n, d = TV.dimension(t), scale, min_entry = lo, max_sum_error = sumerr,
       max_roundtrip = rt, failed_inverses = failed)
end

slate_table([simplex_report(n; scale) for scale in (1.0, 3.0) for n in (3, 10, 50, 200)];
            format = (min_entry = :scientific, max_sum_error = :scientific,
                      max_roundtrip = :scientific))
ndscalemin_entrymax_sum_errormax_roundtripfailed_inverses
3219.76e-32.22e-162.22e-150
10915.35e-44.44e-169.10e-140
504912.24e-56.66e-163.65e-120
20019913.13e-61.33e-154.61e-110
3234.37e-72.22e-161.07e-120
10931.50e-113.33e-161.64e-70
504931.17e-204.44e-162.74e03
20019933.69e-358.88e-164.28e1743

The same thing over Bijectors.jl

Bijectors is organised around distributions rather than around constraints. You ask for the bijector of a distribution and get the map from its support to $\mathbb{R}^k$, so each constraint has to be named by a distribution that carries it.

function probe(d)
    b = BJ.bijector(d)
    v = rand(Xoshiro(1), d)
    u = BJ.transform(b, v)
    (; distribution  = nameof(typeof(d)),
       bijector      = nameof(typeof(b)),
       constrained   = summary(v),
       unconstrained = summary(u))
end

[probe(d) for d in (Exponential(), Uniform(0, 1), Dirichlet(ones(3)),
                    LKJ(4, 1.0), LKJCholesky(4, 1.0))]
5-element Vector{@NamedTuple{distribution::Symbol, bijector::Symbol, constrained::String, unconstrained::String}}:
 (distribution = :Exponential, bijector = :Fix, constrained = "Float64", unconstrained = "Float64")
 (distribution = :Uniform, bijector = :TruncatedBijector, constrained = "Float64", unconstrained = "Float64")
 (distribution = :Dirichlet, bijector = :SimplexBijector, constrained = "3-element Vector{Float64}", unconstrained = "2-element Vector{Float64}")
 (distribution = :LKJ, bijector = :VecCorrBijector, constrained = "4×4 Matrix{Float64}", unconstrained = "6-element Vector{Float64}")
 (distribution = :LKJCholesky, bijector = :VecCholeskyBijector, constrained = "LinearAlgebra.Cholesky{Float64, Matrix{Float64}}", unconstrained = "6-element Vector{Float64}")
# Constraint -> the Bijectors map from the constrained set to ℝᵏ.
bj(::Unbounded)   = identity
bj(::Positive)    = BJ.bijector(Exponential())
bj(::Negative)    = BJ.bijector(Exponential())  BJ.Scale(-1.0)
bj(c::Interval)   = BJ.bijector(Uniform(c.lo, c.hi))
bj(c::Simplex)    = BJ.bijector(Dirichlet(ones(c.n)))
bj(c::CorrMatrix) = BJ.bijector(LKJ(c.n, 1.0))
bj(c::CorrFactor) = BJ.bijector(LKJCholesky(c.n, 1.0))

# Bijectors has no `dimension`, so the sizes are written out.
bj_dim(::Union{Unbounded,Positive,Negative,Interval}) = 1
bj_dim(c::Simplex)    = c.n - 1
bj_dim(c::CorrMatrix) = c.n * (c.n - 1) ÷ 2
bj_dim(c::CorrFactor) = c.n * (c.n - 1) ÷ 2
bj_dim(c::Nested)     = bj_dimension(c.T)

bj_dimension(::Type{S}) where {S} = sum(bj_dim, values(constraints_of(S)))

_bj_inv(c::Constraint, seg) = BJ.with_logabsdet_jacobian(BJ.inverse(bj(c)), seg)
function _bj_inv(c::CorrFactor, seg)
    ch, δ = BJ.with_logabsdet_jacobian(BJ.inverse(bj(c)), seg)
    UpperTriangular(ch.U), δ
end
_bj_inv(c::Nested, seg) = bj_constrain_and_logjac(c.T, seg)

_asvec(v::Number) = [v]
_asvec(v)         = vec(v)
_bj_fwd(c::Constraint, v) = _asvec(BJ.transform(bj(c), v))
_bj_fwd(c::CorrFactor, U) = _asvec(BJ.transform(bj(c), Cholesky(Matrix(U), 'U', 0)))
_bj_fwd(::Nested, v)      = bj_unconstrain(v)

"ℝᵖ -> S, with the log-Jacobian, keeping the offsets by hand."
function bj_constrain_and_logjac(::Type{S}, x::AbstractVector) where {S}
    cs = constraints_of(S)
    i = 1= zero(float(eltype(x)))
    vals = Any[]
    for c in values(cs)
        d = bj_dim(c)
        seg = d == 1 ? x[i] : x[i:(i + d - 1)]
        i += d
        v, δ = _bj_inv(c, seg)
        ℓ += δ
        push!(vals, v)
    end
    _key(S)(; NamedTuple{keys(cs)}(Tuple(vals))...), ℓ
end

bj_constrain(::Type{S}, x::AbstractVector) where {S} = first(bj_constrain_and_logjac(S, x))
bj_unconstrain(s::S) where {S} =
    reduce(vcat, [_bj_fwd(c, getfield(s, k)) for (k, c) in pairs(constraints_of(S))])

nothing
let θb = bj_constrain(S, x)
    (; dimension_agrees = bj_dimension(S) == intrinsic_dimension(S),
       max_value_gap    = maximum(abs, freecoords(θb) .- freecoords(θ)),
       roundtrip_error  = maximum(abs, bj_unconstrain(θb) .- x))
end
(dimension_agrees = true, max_value_gap = 1.1102230246251565e-16, roundtrip_error = 4.440892098500626e-16)

The values agree to machine precision, so the two packages use the same parameterisation of the correlation matrix. The log-Jacobians agree too, both reporting $x \mapsto C$: Bijectors because VecCorrBijector maps to the correlation matrix directly, TransformVariables because tv(::CorrMatrix) adds the factor-to-matrix term.

let (_, ℓb) = bj_constrain_and_logjac(S, x), (_, ℓt) = constrain_and_logjac(S, x)
    ν = numeric_logjac(S, x)
    (; bijectors          = ℓb,
       transformvariables = ℓt,
       numeric            = ν,
       bijectors_error    = ℓb - ν,
       transformvariables_error = ℓt - ν)
end
(bijectors = -13.371643795260162, transformvariables = -13.37164379526016, numeric = -13.371643795260162, bijectors_error = 0.0, transformvariables_error = 1.7763568394002505e-15)
let mb = bj_constrain(M, xm)
    (; dimension       = bj_dimension(M),
       agrees_with_tv  = maximum(abs, bj_unconstrain(mb) .- unconstrain(m)),
       roundtrip_error = maximum(abs, bj_unconstrain(mb) .- xm))
end
(dimension = 10, agrees_with_tv = 1.1102230246251565e-15, roundtrip_error = 1.7763568394002505e-15)

Why you want the mapping

A mixture of $N$ zero-mean Gaussians in $P$ dimensions that share a correlation structure and differ only in scale. The parameters are the three constraints from the thread, and the optimizer works on $x \in \mathbb{R}^9$ without ever seeing one of them.

This fit uses FactorParams, the version that stores the Cholesky factor. Written against MixtureParams it throws PosDefException partway through the line search. The likelihood has to factorize C to evaluate it, the transform computed that factor and discarded it, and near the edge of the parameter space the round trip through a dense C is no longer positive definite. Keeping the factor takes the factorization out of the inner loop, and the two fields parameterise the same nine numbers either way.

C_true = [0.6^abs(i - j) for i in 1:4, j in 1:4]
θ_true = FactorParams(0.7, [0.5, 0.3, 0.2], cholesky(Symmetric(C_true)).U)

function simulate(θ, n; rng = Xoshiro(2026))
    P, N = size.U, 1), length.w)
    edges = cumsum.w)
    Y = zeros(P, n)
    for i in 1:n
        k = something(findfirst((rand(rng)), edges), N)
        Y[:, i] =.σ * k) .*.U' * randn(rng, P))
    end
    Y
end

Y = simulate(θ_true, 4000)
(; truth = θ_true, data = summary(Y))
(truth = Main.NB.FactorParams{3, 4, Float64}(0.7, [0.5, 0.3, 0.2], [1.0 0.6 0.36 0.216; 0.0 0.8 0.48 0.288; 0.0 0.0 0.8 0.48; 0.0 0.0 0.0 0.8]), data = "4×4000 Matrix{Float64}")
function logsumexp(v)
    m = maximum(v)
    m + log(sum(x -> exp(x - m), v))
end

"log N(y; 0, Σ) from the factor U of Σ = U'U, without factorizing anything."
function mvn_logpdf(U, y)
    z = U' \ y
    -(length(y) * log(2π) + 2 * sum(log, diag(U)) + sum(abs2, z)) / 2
end

function loglik(θ, Y)
    # A line search can walk to a correlation of ±1 or a scale that underflows, where the
    # scaled factor is singular. Return -Inf there instead of throwing.
    θ.σ * minimum(diag.U)) > 0 || return convert(typeof.σ), -Inf)
    N, logw = length.w), log..w)
    Us = [(θ.σ * k) .* θ.U for k in 1:N]
    sum(1:size(Y, 2)) do i
        y = @view Y[:, i]
        logsumexp([logw[k] + mvn_logpdf(Us[k], y) for k in 1:N])
    end
end

loglik(θ_true, Y)
-22823.640085784325
nll(x)     = -loglik(constrain(FactorParams{3,4}, x), Y)
nll!(G, x) = ForwardDiff.gradient!(G, nll, x)

x₀  = zeros(intrinsic_dimension(FactorParams{3,4}))
res = Optim.optimize(nll, nll!, x₀,
                     Optim.LBFGS(; linesearch = LineSearches.BackTracking()))
θ̂   = constrain(FactorParams{3,4}, Optim.minimizer(res))
Ĉ   = Matrix(Symmetric(θ̂.U'θ̂.U))

(; iterations  = Optim.iterations(res),
   converged   = Optim.converged(res),
   loglik_fit  = -Optim.minimum(res),
   loglik_true = loglik(θ_true, Y),
   σ           = (fit = round(θ̂.σ; digits = 4), truth = θ_true.σ),
   w           = (fit = round.(θ̂.w; digits = 4), truth = θ_true.w),
   max_C_error = maximum(abs, Ĉ .- C_true))
(iterations = 31, converged = true, loglik_fit = -22820.945331823918, loglik_true = -22823.640085784325, σ = (fit = 0.7018, truth = 0.7), w = (fit = [0.5109, 0.2995, 0.1896], truth = [0.5, 0.3, 0.2]), max_C_error = 0.018823611013218833)
let ls = 0:3
    pts = [(l, Ĉ[i, j]) for l in ls for i in 1:4 for j in 1:4 if j - i == l]
    echart(series(:line, collect(ls), [0.6^l for l in ls];
                  name = "truth, 0.6ᵏ", symbol = "none", lineStyle = (type = :dashed,)),
           series(:scatter, first.(pts), last.(pts); name = "recovered", symbolSize = 11);
           title     = "Correlation by lag, recovered from 4000 draws",
           xAxis     = (name = "lag |i − j|", interval = 1),
           yAxis     = (name = "correlation", min = 0, max = 1.05),
           legend    = true,
           animation = false,
           height    = 320)
end

What this comes to

Most of what the thread asks for already exists.

The bijections and their log-Jacobians are in both packages, and they agree to machine precision on every value computed above.

The binding to a struct type is mostly there as well, in TransformVariables.as(::Type, inner). It asks the struct for a ConstructionBase.constructorof that keeps the type parameters, and for construction by keyword.

What is left is the macro. It writes the struct, the strict constructor, and constraints_of, which is all either backend reads. Adding a constraint is a new descriptor plus one method per backend.

Three of the decisions inside it were not obvious beforehand.

Sizes go in type parameters and the element type goes last. intrinsic_dimension is then a function of the type, which is what makes nesting compose, and Name{sizes...} lets a transform build the struct at ForwardDiff.Dual.

The strict constructor and the transform are separate paths. The transform's image is already in the constraint set, and re-checking it fails at the floating-point edge: a line search reaches exp(x) == 0.0 and the strict constructor throws.

A correlation matrix field is a second change of variables on top of the Cholesky factor. corr_cholesky_factor reports $x \mapsto U$, and a field holding $C = U'U$ needs $\sum_i (n-i)\log U_{ii}$ added to it. It also puts a Cholesky in every likelihood evaluation, which corrfactor avoids.

Which package

TransformVariables fits this shape better. It has dimension, so the offsets into x are bookkeeping it already does, and as(::Type, inner) is the struct binding. Bijectors is organised around distributions, so each constraint has to be named by a distribution that carries it, and the offsets are yours to track. Its VecCorrBijector does map to the correlation matrix directly, so it reports $x \mapsto C$ with no extra term.

Run this notebook live

Get the full interactive notebook (with the AI agent) on your machine. Needs Julia 1.10+. The launch script installs Kaimon + KaimonSlate and starts the notebook (its exact environment is reconstructed from the bundle).

One-lineryour platform (paste into a terminal):

Show the command for another OS ▾

Or download run.jl to inspect first, then julia run.jl. On Windows, grab run.bat + run.ps1 beside it and double-click run.bat. Just the env? download the bundle.