From 43c06abe6ddf71355ec535a7ec315e0c4de103a1 Mon Sep 17 00:00:00 2001 From: David Xu <42751767+zuhengxu@users.noreply.github.com> Date: Wed, 23 Aug 2023 12:33:44 -0700 Subject: [PATCH] add more synthetic targets (#20) * add Neal's Funnel and Warped Gaussian * fixed bug in warped gaussian * add reference for warped Gauss * add Cross dsitribution * udpate docs for cross * Update example/targets/cross.jl change comment into docs Co-authored-by: Tor Erlend Fjelde * Update example/targets/cross.jl Co-authored-by: Tor Erlend Fjelde * update cross docs * minor ed * Update example/targets/neal_funnel.jl Co-authored-by: Tor Erlend Fjelde * Update example/targets/neal_funnel.jl Co-authored-by: Tor Erlend Fjelde * doc banana using * fixing docs with latex * baanan docs with latex * add NF quick intro * Revert "add NF quick intro" This reverts commit e399274450946207c3e38a1fbd23f7cc97b7a486. * rm unnecesary code for cross * rm example/manifest * Update example/targets/neal_funnel.jl Co-authored-by: Tor Erlend Fjelde * Update example/targets/cross.jl Co-authored-by: Tor Erlend Fjelde * Update example/targets/cross.jl Co-authored-by: Tor Erlend Fjelde * minor update to cross docs --------- Co-authored-by: Tor Erlend Fjelde --- example/targets/banana.jl | 22 +++++--- example/targets/cross.jl | 39 +++++++++++++ example/targets/neal_funnel.jl | 63 +++++++++++++++++++++ example/targets/warped_gaussian.jl | 90 ++++++++++++++++++++++++++++++ 4 files changed, 206 insertions(+), 8 deletions(-) create mode 100644 example/targets/cross.jl create mode 100644 example/targets/neal_funnel.jl create mode 100644 example/targets/warped_gaussian.jl diff --git a/example/targets/banana.jl b/example/targets/banana.jl index b72b555..ee89cb5 100644 --- a/example/targets/banana.jl +++ b/example/targets/banana.jl @@ -8,22 +8,25 @@ using IrrationalConstants Multidimensional banana-shape distribution. # Fields -- 'dim::Int': Dimension of the distribution, must be >= 2 -- 'b::T': Banananicity constant, the larger "|b|" the more curved the banana -- 'var::T': Variance of the first dimension, must be > 0 - +$(FIELDS) # Explanation The banana distribution is obtained by applying a transformation ϕ to a multivariate normal -distribution "N(0, diag(var, 1, 1, …, 1))". The transformation ϕ is defined as -"ϕ(x₁, … , xₚ) = (x₁, x₂ - B x₁² + var*B, x₃, … , xₚ)", +distribution ``\\mathcal{N}(0, \\text{diag}(var, 1, 1, …, 1))``. The transformation ϕ is defined as +```math +\phi(x_1, … , x_p) = (x_1, x_2 - B x_1^² + \text{var}*B, x_3, … , x_p) +```` which has a unit Jacobian determinant. Hence the density "fb" of a p-dimensional banana distribution is given by -"fb(x₁, … , xₚ) = exp[ -½x₁²/var - ½(x₂ + B x₁² - var*B)² - ½(x₃² + x₄² + … + xₚ²)] / Z", +```math +fb(x_1, \dots, x_p) = \exp\left[ -\frac{1}{2}\frac{x_1^2}{\text{var}} - +\frac{1}{2}(x_2 + B x_1^2 - \text{var}*B)^2 - \frac{1}{2}(x_3^2 + x_4^2 + \dots ++ x_p^2) \right] / Z, +``` where "B" is the "banananicity" constant, determining the curvature of a banana, and -"Z = sqrt(var * (2π)^p))" is the normalization constant. +``Z = \\sqrt{\\text{var} * (2\\pi)^p)}`` is the normalization constant. # Reference @@ -33,8 +36,11 @@ Gareth O. Roberts and Jeffrey S. Rosenthal Journal of computational and graphical statistics, Volume 18, Number 2 (2009): 349-367. """ struct Banana{T<:Real} <: ContinuousMultivariateDistribution + "Dimension of the distribution, must be >= 2" dim::Int # Dimension + "Banananicity constant, the larger |b| the more curved the banana" b::T # Curvature + "Variance of the first dimension, must be > 0" var::T # Variance function Banana{T}(dim::Int, b::T, var::T) where {T<:Real} dim >= 2 || error("dim must be >= 2") diff --git a/example/targets/cross.jl b/example/targets/cross.jl new file mode 100644 index 0000000..7c843f8 --- /dev/null +++ b/example/targets/cross.jl @@ -0,0 +1,39 @@ +using Distributions, Random +""" + Cross(μ::Real=2.0, σ::Real=0.15) + +2-dimensional Cross distribution + + +# Explanation + +The Cross distribution is a 2-dimension 4-component Gaussian distribution with a "cross" +shape that is symmetric about the y- and x-axises. The mixture is defined as + +```math +\begin{aligned} +p(x) = +& 0.25 \mathcal{N}(x | (0, \mu), (\sigma, 1)) + \\ +& 0.25 \mathcal{N}(x | (\mu, 0), (1, \sigma)) + \\ +& 0.25 \mathcal{N}(x | (0, -\mu), (\sigma, 1)) + \\ +& 0.25 \mathcal{N}(x | (-\mu, 0), (1, \sigma))) +\end{aligned} +``` + +where ``μ`` and ``σ`` are the mean and standard deviation of the Gaussian components, +respectively. See an example of the Cross distribution in Page 18 of [1]. + +# Reference +[1] Zuheng Xu, Naitong Chen, Trevor Campbell +"MixFlows: principled variational inference via mixed flows." +International Conference on Machine Learning, 2023 +""" +Cross() = Cross(2.0, 0.15) +function Cross(μ::T, σ::T) where {T<:Real} + return MixtureModel([ + MvNormal([zero(μ), μ], [σ, one(σ)]), + MvNormal([-μ, one(μ)], [one(σ), σ]), + MvNormal([μ, one(μ)], [one(σ), σ]), + MvNormal([zero(μ), -μ], [σ, one(σ)]), + ]) +end diff --git a/example/targets/neal_funnel.jl b/example/targets/neal_funnel.jl new file mode 100644 index 0000000..66f58b0 --- /dev/null +++ b/example/targets/neal_funnel.jl @@ -0,0 +1,63 @@ +using Distributions, Random + +""" + Funnel{T<:Real} + +Multidimensional Neal's Funnel distribution + +# Fields +$(FIELDS) + +# Explanation + +The Neal's Funnel distribution is a p-dimensional distribution with a funnel shape, +originally proposed by Radford Neal in [2]. +The marginal distribution of ``x_1`` is Gaussian with mean "μ" and standard +deviation "σ". The conditional distribution of ``x_2, \dots, x_p | x_1`` are independent +Gaussian distributions with mean 0 and standard deviation ``\\exp(x_1/2)``. +The generative process is given by +```math +x_1 \sim \mathcal{N}(\mu, \sigma^2), \quad x_2, \ldots, x_p \sim \mathcal{N}(0, \exp(x_1)) +``` + + +# Reference +[1] Stan User’s Guide: +https://mc-stan.org/docs/2_18/stan-users-guide/reparameterization-section.html#ref-Neal:2003 +[2] Radford Neal 2003. “Slice Sampling.” Annals of Statistics 31 (3): 705–67. +""" +struct Funnel{T<:Real} <: ContinuousMultivariateDistribution + "Dimension of the distribution, must be >= 2" + dim::Int + "Mean of the first dimension" + μ::T + "Standard deviation of the first dimension, must be > 0" + σ::T + function Funnel{T}(dim::Int, μ::T, σ::T) where {T<:Real} + dim >= 2 || error("dim must be >= 2") + σ > 0 || error("σ must be > 0") + return new{T}(dim, μ, σ) + end +end +Funnel(dim::Int, μ::T, σ::T) where {T<:Real} = Funnel{T}(dim, μ, σ) +Funnel(dim::Int, σ::T) where {T<:Real} = Funnel{T}(dim, zero(T), σ) +Funnel(dim::Int) = Funnel(dim, 0.0, 9.0) + +Base.length(p::Funnel) = p.dim +Base.eltype(p::Funnel{T}) where {T<:Real} = T + +function Distributions._rand!(rng::AbstractRNG, p::Funnel, x::AbstractVecOrMat) + T = eltype(x) + d, μ, σ = p.dim, p.μ, p.σ + d == size(x, 1) || error("Dimension mismatch") + x[1, :] .= randn(rng, T, size(x, 2)) .* σ .+ μ + x[2:end, :] .= randn(rng, T, d - 1, size(x, 2)) .* exp.(@view(x[1, :]) ./ 2)' + return x +end + +function Distributions._logpdf(p::Funnel, x::AbstractVector) + d, μ, σ = p.dim, p.μ, p.σ + lpdf1 = logpdf(Normal(μ, σ), x[1]) + lpdfs = logpdf.(Normal.(zeros(T, d - 1), exp(x[1] / 2)), @view(x[2:end])) + return lpdf1 + sum(lpdfs) +end diff --git a/example/targets/warped_gaussian.jl b/example/targets/warped_gaussian.jl new file mode 100644 index 0000000..2179046 --- /dev/null +++ b/example/targets/warped_gaussian.jl @@ -0,0 +1,90 @@ +using Distributions, Random, LinearAlgebra, IrrationalConstants + +""" + WarpedGauss{T<:Real} + +2-dimensional warped Gaussian distribution + +# Fields +$(FIELDS) + +# Explanation + +The banana distribution is obtained by applying a transformation ϕ to a 2-dimensional normal +distribution ``\\mathcal{N}(0, diag(\\sigma_1, \\sigma_2))``. The transformation ϕ(x) is defined as +```math +ϕ(x_1, x_2) = (r*\cos(\theta + r/2), r*\sin(\theta + r/2)), +``` +where ``r = \\sqrt{x\_1^2 + x_2^2}``, ``\\theta = \\atan(x₂, x₁)``, +and "atan(y, x) ∈ [-π, π]" is the angle, in radians, between the positive x axis and the +ray to the point "(x, y)". See page 18. of [1] for reference. + + +# Reference +[1] Zuheng Xu, Naitong Chen, Trevor Campbell +"MixFlows: principled variational inference via mixed flows." +International Conference on Machine Learning, 2023 +""" +struct WarpedGauss{T<:Real} <: ContinuousMultivariateDistribution + "Standard deviation of the first dimension, must be > 0" + σ1::T + "Standard deviation of the second dimension, must be > 0" + σ2::T + function WarpedGauss{T}(σ1, σ2) where {T<:Real} + σ1 > 0 || error("σ₁ must be > 0") + σ2 > 0 || error("σ₂ must be > 0") + return new{T}(σ1, σ2) + end +end +WarpedGauss(σ1::T, σ2::T) where {T<:Real} = WarpedGauss{T}(σ1, σ2) +WarpedGauss() = WarpedGauss(1.0, 0.12) + +Base.length(p::WarpedGauss) = 2 +Base.eltype(p::WarpedGauss{T}) where {T<:Real} = T +Distributions.sampler(p::WarpedGauss) = p + +# Define the transformation function φ and the inverse ϕ⁻¹ for the warped Gaussian distribution +function ϕ!(p::WarpedGauss, z::AbstractVector) + length(z) == 2 || error("Dimension mismatch") + x, y = z + r = norm(z) + θ = atan(y, x) #in [-π , π] + θ -= r / 2 + z .= r .* [cos(θ), sin(θ)] + return z +end + +function ϕ⁻¹(p::WarpedGauss, z::AbstractVector) + length(z) == 2 || error("Dimension mismatch") + x, y = z + r = norm(z) + θ = atan(y, x) #in [-π , π] + # increase θ depending on r to "smear" + θ += r / 2 + + # get the x,y coordinates foαtransformed point + xn = r * cos(θ) + yn = r * sin(θ) + # compute jacobian + logJ = log(r) + return [xn, yn], logJ +end + +function Distributions._rand!(rng::AbstractRNG, p::WarpedGauss, x::AbstractVecOrMat) + size(x, 1) == 2 || error("Dimension mismatch") + σ₁, σ₂ = p.σ₁, p.σ₂ + randn!(rng, x) + x .*= [σ₁, σ₂] + for y in eachcol(x) + ϕ!(p, y) + end + return x +end + +function Distributions._logpdf(p::WarpedGauss, x::AbstractVector) + size(x, 1) == 2 || error("Dimension mismatch") + σ₁, σ₂ = p.σ₁, p.σ₂ + S = [σ₁, σ₂] .^ 2 + z, logJ = ϕ⁻¹(p, x) + return -sum(z .^ 2 ./ S) / 2 - IrrationalConstants.log2π - log(σ₁) - log(σ₂) + logJ +end