Constrained structs with automatic parameter handling
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))"
nothingWhat a constraint is
Each descriptor names three things:
- a subset $\mathcal{S}$ of a vector space, tested by
satisfies - its dimension $d$ as a smooth manifold
- 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)
nothingThe 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)
endMain.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
"ν" => 1The 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
end5-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,))| constraint | d | in_set | left_inverse | right_inverse | logjac |
|---|---|---|---|---|---|
| positive | 1 | true | 1.11e-16 | 0e0 | 1.11e-16 |
| negative | 1 | true | 1.11e-16 | 0e0 | 1.11e-16 |
| in (-2.0, 3.0) | 1 | true | 5.77e-15 | 8.88e-16 | 1.33e-15 |
| on the 5-element unit simplex | 4 | true | 1.09e-13 | 1.67e-16 | 3.55e-15 |
| a 4×4 correlation Cholesky factor | 6 | true | 3.88e-11 | 7.22e-16 | 2.42e-13 |
| a 4×4 correlation matrix | 6 | true | 7.04e-11 | 3.44e-15 | 2.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))| n | d | scale | min_entry | max_sum_error | max_roundtrip | failed_inverses |
|---|---|---|---|---|---|---|
| 3 | 2 | 1 | 9.76e-3 | 2.22e-16 | 2.22e-15 | 0 |
| 10 | 9 | 1 | 5.35e-4 | 4.44e-16 | 9.10e-14 | 0 |
| 50 | 49 | 1 | 2.24e-5 | 6.66e-16 | 3.65e-12 | 0 |
| 200 | 199 | 1 | 3.13e-6 | 1.33e-15 | 4.61e-11 | 0 |
| 3 | 2 | 3 | 4.37e-7 | 2.22e-16 | 1.07e-12 | 0 |
| 10 | 9 | 3 | 1.50e-11 | 3.33e-16 | 1.64e-7 | 0 |
| 50 | 49 | 3 | 1.17e-20 | 4.44e-16 | 2.74e0 | 3 |
| 200 | 199 | 3 | 3.69e-35 | 8.88e-16 | 4.28e1 | 743 |
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))])
nothinglet θ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)
endWhat 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.