From 357621126a8eacb7bb6c412a0c0fc8ff062f047f Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 22 Jan 2025 19:00:35 +0100 Subject: [PATCH 01/12] fix: implement base normal gamma manifold --- Project.toml | 8 ++ src/ExponentialFamilyManifolds.jl | 3 + src/normal_gamma/natural.jl | 113 ++++++++++++++++++++++ src/normal_gamma/partial.jl | 11 +++ test/normal_gamma_tests.jl | 156 ++++++++++++++++++++++++++++++ 5 files changed, 291 insertions(+) create mode 100644 src/normal_gamma/natural.jl create mode 100644 src/normal_gamma/partial.jl create mode 100644 test/normal_gamma_tests.jl diff --git a/Project.toml b/Project.toml index f7e535c..16f2892 100644 --- a/Project.toml +++ b/Project.toml @@ -6,22 +6,30 @@ version = "1.5.0" [deps] BayesBase = "b4ee3484-f114-42fe-b91c-797d54a0c67e" ExponentialFamily = "62312e5e-252a-4322-ace9-a5f4bf9b357b" +FastCholesky = "2d5283b6-8564-42b6-bb00-83ed8e915756" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e" ManifoldsBase = "3362f125-f0bb-47a3-aa74-596ffd7ef2fb" +NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" +SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [compat] BayesBase = "1.3" ExponentialFamily = "1.6.0" +FastCholesky = "1.3.1" LinearAlgebra = "1.10" Manifolds = "0.9" ManifoldsBase = "0.15" +NLsolve = "4.5.1" Random = "1.10" RecursiveArrayTools = "3" +SpecialFunctions = "2.5.0" Static = "0.8, 1" +StaticArrays = "1.9.10" julia = "1.10" [extras] diff --git a/src/ExponentialFamilyManifolds.jl b/src/ExponentialFamilyManifolds.jl index 5a38fa0..1dd76b6 100644 --- a/src/ExponentialFamilyManifolds.jl +++ b/src/ExponentialFamilyManifolds.jl @@ -27,4 +27,7 @@ include("natural_manifolds/poisson.jl") include("natural_manifolds/rayleigh.jl") include("natural_manifolds/weibull.jl") include("natural_manifolds/wishart.jl") +# normal gamma +include("normal_gamma/partial.jl") +include("normal_gamma/natural.jl") end diff --git a/src/normal_gamma/natural.jl b/src/normal_gamma/natural.jl new file mode 100644 index 0000000..c86c71c --- /dev/null +++ b/src/normal_gamma/natural.jl @@ -0,0 +1,113 @@ +using ExponentialFamily +using FastCholesky +using LinearAlgebra + +import ManifoldsBase: check_point, check_vector, manifold_dimension, representation_size, exp!, is_point, is_vector +import ManifoldsBase: zero_vector, zero_vector!, default_retraction_method +import ManifoldsBase: ExponentialRetraction +import ManifoldsBase: injectivity_radius + +""" +Compute all Christoffel symbols Γ^k_{i j}(x). +Returns a 4×4×4 array: Gamma[k,i,j]. +""" +function christoffel(::Type{NormalGamma}, ::NaturalParametersSpace, η) + ef = ExponentialFamilyDistribution(NormalGamma, η, nothing) + G = fisherinformation(ef) + G⁻¹ = cholinv(G) + dG = partial_metric(NormalGamma, NaturalParametersSpace(), η) + # formula: Γ^k_{i,j} = 1/2 * sum_{ℓ} g^{kℓ} ( ∂i g_{ℓj} + ∂j g_{ℓi} - ∂ℓ g_{i j} ) + Γ = zeros(4,4,4) + for k in 1:4 + for i in 1:4 + for j in 1:4 + s = 0.0 + for ℓ in 1:4 + s += G⁻¹[k,ℓ] * (dG[ℓ,j,i] + dG[ℓ,i,j] - dG[i,j,ℓ]) + end + Γ[k,i,j] = 0.5 * s + end + end + end + return Γ +end + +function exp_secondorder( + ::Type{NormalGamma}, + ::NaturalParametersSpace, + p0, + v0 + ) + Γ = christoffel(NormalGamma, NaturalParametersSpace(), p0) + + Δ = zeros(4) + for k in 1:4 + for i in 1:4 + for j in 1:4 + Δ[k] += Γ[k,i,j] * v0[i] * v0[j] + end + end + Δ[k] *= -0.5 + end + + return p0 + v0 + Δ + end + + +""" + NormalGammaNaturalManifold <: AbstractManifold{ℝ} + +4D manifold for NormalGamma in natural params η = (η1, η2, η3, η4). +Domain constraints, etc. +""" +struct NormalGammaNaturalManifold <: AbstractManifold{ℝ} end + +function manifold_dimension(::NormalGammaNaturalManifold) + return 4 +end + +function representation_size(::NormalGammaNaturalManifold) + return (4,) +end + +function check_point(::NormalGammaNaturalManifold, η) + cond_isproper = isproper(NaturalParametersSpace(), NormalGamma, η, nothing) + if !cond_isproper + return DomainError("$η is not a valid point on NormalGammaNaturalManifold") + end + return nothing +end + +function check_vector(::NormalGammaNaturalManifold, η, X) + if length(X)!=4 + return DomainError(X, "Vector must be length 4.") + end + return nothing +end + +function exp!(::NormalGammaNaturalManifold, η_out::AbstractVector, η_in::AbstractVector, X_in::AbstractVector, t::Real=1.0) + η_out .= exp_secondorder(NormalGamma, NaturalParametersSpace(), η_in, t.*X_in) + return η_out +end + +# Provide a zero_vector etc. +function zero_vector(::NormalGammaNaturalManifold, η) + return zeros(4) +end + +function zero_vector!(::NormalGammaNaturalManifold, X, η) + fill!(X, 0.0) + return X +end + +default_retraction_method(::NormalGammaNaturalManifold) = ExponentialRetraction() + +injectivity_radius(::NormalGammaNaturalManifold) = Inf + +import ManifoldsBase: inner + +function inner(::NormalGammaNaturalManifold, p, X, Y) + ef = ExponentialFamilyDistribution(NormalGamma, p, nothing) + G = fisherinformation(ef) + return dot(X, G, Y) +end \ No newline at end of file diff --git a/src/normal_gamma/partial.jl b/src/normal_gamma/partial.jl new file mode 100644 index 0000000..2ee85e2 --- /dev/null +++ b/src/normal_gamma/partial.jl @@ -0,0 +1,11 @@ +using SpecialFunctions +using ExponentialFamily + +function partial_metric_tuple(::Type{NormalGamma}, ::NaturalParametersSpace, η1, η2, η3, η4) + ((+)((/)((*)((*)(-8.0, (/)((*)(2, (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2)))), η1), 4.0), (*)(0.25, (+)((/)((*)((*)(-16.0, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (*)((*)(-8.0, (+)((/)((*)(2, (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (/)((*)((*)((*)((*)(-64.0, (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η1), η2), 4.0))), η1)))), (+)((/)((*)((*)((*)(-2, (+)((/)((*)(8.0, η1), (*)(4.0, η2)), (/)((*)(-32.0, η1), (*)(16.0, η2)))), η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (*)((+)((*)(-4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (*)((*)(16.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), η2)), (+)((/)((*)(2, (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (/)((*)((*)((*)((*)(-64.0, (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η1), η2), 4.0)))), (+)((/)(-2, (*)((*)(4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η2)), (/)((*)((*)(-8.0, (/)((*)(-2, η1), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2)))), η1), 4.0)), (*)((*)(4.0, (+)((/)((*)(2, (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (/)((*)((*)((*)((*)(-64.0, (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η1), η2), 4.0))), η2), (+)((/)((*)((*)((*)(-2, (+)((/)((*)(8.0, η1), (*)(4.0, η2)), (/)((*)(-32.0, η1), (*)(16.0, η2)))), η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (*)((+)((*)(-4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (*)((*)(16.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), η2)), (+)((/)((*)(2, (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (/)((*)((*)((*)((*)(-64.0, (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η1), η2), 4.0)))), (+)((+)((+)((/)((*)((*)(-2, (/)((*)((*)((*)(128.0, (/)((^)(η1, 2), (*)(256.0, (^)(η2, 4)))), η2), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2))), η1), (*)(4.0, η2)), (/)((*)((*)(4.0, (+)((/)((*)((*)((*)(-4, (/)((*)((*)(-4.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η1), (*)(4.0, η2)), (/)((*)((*)(-8.0, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))))), (^)(η1, 2)), (*)(16.0, (^)(η2, 2)))), (/)((*)((*)(256.0, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η2, 3)))), (/)((*)((*)((*)(-32.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2)))), (+)((/)((*)(8.0, η1), (*)((*)(16.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η2, 2))), (/)((*)((*)((*)(-32.0, (/)((*)(4.0, (^)(η1, 2)), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 4)))), η1), η2), 4.0)), (+)((/)((*)((*)((*)(-4, (/)((*)((*)(-4.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η1), (*)(4.0, η2)), (/)((*)((*)(-8.0, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2)))), (+)((/)(-2, (*)((*)(4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η2)), (/)((*)((*)(-8.0, (/)((*)(-2, η1), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2)))), η1), 4.0)), (+)((/)((*)(8.0, η1), (*)((*)(16.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η2, 2))), (/)((*)((*)((*)(-32.0, (/)((*)(4.0, (^)(η1, 2)), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 4)))), η1), η2), 4.0)), 0, (/)((*)((*)(-2, (/)(1, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2))), η1), (*)(4.0, η2)), (*)((*)(4.0, (+)((/)((*)(2, (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (/)((*)((*)((*)((*)(-64.0, (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η1), η2), 4.0))), η2), (+)((/)((*)((*)((*)(-4, (/)((*)((*)(-4.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η1), (*)(4.0, η2)), (/)((*)((*)(-8.0, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2)))), (/)((*)((*)(-2, (/)(1, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2))), η1), (*)(4.0, η2)), (/)((*)((*)((*)(-4, (/)((+)(0.5, η3), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η1), (*)(4.0, η2)), (+)((*)((*)(-1, (+)((*)(4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (*)((*)(-16.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), η2))), (/)((*)(2, (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2)))), (*)((*)((*)(2.0, (+)((*)((*)(32.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), η2), (*)((*)((*)(-128.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η2, 2)))), (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4)))), η1)), (+)((/)((*)((*)((*)(-2, (+)((+)((/)((*)(-16.0, (^)(η1, 2)), (*)(16.0, (^)(η2, 2))), (/)((*)(-16.0, (^)(η1, 2)), (*)(16.0, (^)(η2, 2)))), (*)((*)(512.0, (/)((^)(η1, 2), (*)(256.0, (^)(η2, 4)))), (^)(η2, 2)))), η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (*)((*)((*)(-1, (+)((*)((*)(32.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), η2), (*)((*)((*)(-128.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η2, 2)))), (+)((*)(-4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (*)((*)(16.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), η2))), (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4))))), (*)((*)(-1, (+)((*)(4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (*)((*)(-16.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), η2))), (/)((*)(-2, η1), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2)))), (+)((/)((*)((*)(8.0, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (*)((*)((*)(-4.0, (+)((*)((*)(32.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), η2), (*)((*)((*)(-128.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η2, 2)))), (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4)))), η2)), (+)((/)((*)((*)((*)(-2, (+)((+)((/)((*)(-16.0, (^)(η1, 2)), (*)(16.0, (^)(η2, 2))), (/)((*)(-16.0, (^)(η1, 2)), (*)(16.0, (^)(η2, 2)))), (*)((*)(512.0, (/)((^)(η1, 2), (*)(256.0, (^)(η2, 4)))), (^)(η2, 2)))), η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (*)((*)((*)(-1, (+)((*)((*)(32.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), η2), (*)((*)((*)(-128.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η2, 2)))), (+)((*)(-4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (*)((*)(16.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), η2))), (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4))))), (+)((+)((+)((+)((/)((+)((/)((*)((*)(128.0, (^)(η1, 2)), (+)(-0.5, (*)(-1, η3))), (*)(256.0, (^)(η2, 4))), (*)((*)((*)(-131072.0, (/)((^)(η1, 2), (*)(65536.0, (^)(η2, 8)))), (^)(η2, 4)), (+)(-0.5, (*)(-1, η3)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (/)((*)((*)(4.0, (+)((/)((*)((*)((*)(128.0, (/)((^)(η1, 2), (*)(256.0, (^)(η2, 4)))), η2), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (*)((*)((*)(8.0, (/)((*)((*)(-4.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4))), (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))))), (^)(η1, 2)), (*)(16.0, (^)(η2, 2)))), (/)((*)((*)((*)((*)(512.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (/)((^)(η1, 2), (*)(256.0, (^)(η2, 4)))), η2), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2))), (*)((*)(4.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (/)((*)((*)((*)(128.0, (/)((^)(η1, 2), (*)(256.0, (^)(η2, 4)))), η2), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)))), (*)((*)(-8.0, (/)(2.0, (*)(16.0, (^)(η2, 4)))), η2)), (*)((*)(-1, (/)((*)(4.0, (^)(η1, 2)), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 4)))), (+)((*)((*)(32.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η2), (*)((*)(-64.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (^)(η2, 2)))), (+)((/)((*)((*)((*)(128.0, (/)((^)(η1, 2), (*)(256.0, (^)(η2, 4)))), η2), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (*)((*)((*)(8.0, (/)((*)((*)(-4.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4))), (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)))), (*)((*)(-1, (+)((*)(4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (*)((*)(-16.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), η2))), (/)((*)(-2, η1), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2)))), (*)((*)(-1, (/)((*)(4.0, (^)(η1, 2)), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 4)))), (+)((*)((*)(32.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η2), (*)((*)(-64.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (^)(η2, 2)))), 0, (*)((*)(4.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (/)(1, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2))), (+)((/)((*)((*)(8.0, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (*)((*)((*)(-4.0, (+)((*)((*)(32.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), η2), (*)((*)((*)(-128.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η2, 2)))), (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4)))), η2)), (+)((/)((*)((*)((*)(128.0, (/)((^)(η1, 2), (*)(256.0, (^)(η2, 4)))), η2), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (*)((*)((*)(8.0, (/)((*)((*)(-4.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4))), (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)))), (*)((*)(4.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (/)(1, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2))), (*)((*)((*)(8.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (/)((+)(0.5, η3), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (+)((/)(-2, (*)((*)(4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η2)), (/)((*)(4.0, (^)(η1, 2)), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2)))), (/)((*)((*)(-2, (+)((*)(-4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (*)((*)(16.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), η2))), η1), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), 0, (/)((*)(-8.0, η1), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), η2)), (/)((*)((*)(-2, (+)((*)(-4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (*)((*)(16.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), η2))), η1), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (+)((/)((*)(16.0, (^)(η1, 4)), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 4))), (/)((*)(-128.0, (^)(η1, 2)), (*)((*)(256.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η2, 3)))), 0, (/)((*)(4.0, (^)(η1, 2)), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), 0, 0, (SpecialFunctions.polygamma)(2, (+)(0.5, η3)), 0, (/)((*)(-8.0, η1), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), η2)), (/)((*)(4.0, (^)(η1, 2)), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), 0, (/)(1, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (+)((*)((*)(4.0, (/)((*)(2, (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2)))), η2), (*)((*)((*)((*)(-64.0, (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), η1), (^)(η2, 2))), (+)((/)((*)((*)(8.0, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (*)((*)((*)((*)(32.0, (+)((*)(-4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (*)((*)(16.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), η2))), (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η2, 2))), (*)((*)(4.0, (/)((*)(-2, η1), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2)))), η2), (*)((*)((*)(128.0, (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η2, 3)), (+)((/)((*)((*)(8.0, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2))), (*)((*)((*)((*)(32.0, (+)((*)(-4.0, (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (*)((*)(16.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), η2))), (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η2, 2))), (+)((/)((*)((*)((*)(8.0, (/)((*)((*)(-4.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η1, 2)), (*)(16.0, (^)(η2, 2))), (/)((*)((*)((*)(128.0, (/)((^)(η1, 2), (*)(256.0, (^)(η2, 4)))), η2), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2))), (*)((*)(16.0, (/)((*)(4.0, (^)(η1, 2)), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 4)))), (^)(η2, 2)), (*)((*)(2, (/)((*)((*)(-4.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (*)((*)(4.0, (/)((*)(-2, η1), (*)((*)(16.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 2)))), η2), (*)((*)(16.0, (/)((*)(4.0, (^)(η1, 2)), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (^)(η2, 4)))), (^)(η2, 2)), 0, (/)(1, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (*)((*)((*)(128.0, (/)((*)((*)(2, η1), (+)(-0.5, (*)(-1, η3))), (*)((*)(256.0, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4)), (^)(η2, 4)))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (^)(η2, 3)), (*)((*)(2, (/)((*)((*)(-4.0, (/)((^)(η1, 2), (*)(16.0, (^)(η2, 2)))), (+)(-0.5, (*)(-1, η3))), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4))), (/)(1, (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 2)), (*)((*)(2, (/)((+)(0.5, η3), (^)((+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)), 4))), (+)((/)((^)(η1, 2), (*)(4.0, η2)), (*)(-1, η4)))) +end + +function partial_metric(T::Type{NormalGamma}, S::NaturalParametersSpace, η) + tup64 = partial_metric_tuple(T, S, η[1], η[2], η[3], η[4]) + return reshape(collect(tup64), 4, 4, 4) +end diff --git a/test/normal_gamma_tests.jl b/test/normal_gamma_tests.jl new file mode 100644 index 0000000..aa211fe --- /dev/null +++ b/test/normal_gamma_tests.jl @@ -0,0 +1,156 @@ +@testitem "Partial derivatives of NormalGamma Fisher metric in natural coords" begin + + import ExponentialFamilyManifolds: partial_metric + using ExponentialFamily + using ForwardDiff + using BenchmarkTools + + # define some test points in domain: + test_points = [ + [1.0, -2.0, 0.5, -1.0], + [1.5, -1.0, 0.3, -0.9], + [2.0, -3.0, 1.0, -2.0], + [0.5, -0.5, 0.1, -0.1], + ] + + d = NormalGamma(1.0, 1.0, 1.0, 1.0) + ef = convert(ExponentialFamilyDistribution, d) + locfisher = (η) -> fisherinformation(ef, η) + + + # A helper that flattens + function fisher_metric_vec(η) + G = locfisher(η) + return vec(G) + end + + function forwarddiff_partials(η) + FD_jac = ForwardDiff.jacobian(fisher_metric_vec, η) + out = Array{Float64}(undef, 4,4,4) + for i in 1:4 + for j in 1:4 + r = 4*(i-1)+j + for c in 1:4 + out[i,j,c] = FD_jac[r,c] + end + end + end + return out + end + + for η in test_points + fd_partials = forwarddiff_partials(η) + sym_partials = partial_metric(NormalGamma, NaturalParametersSpace(), η) + + diff_array = fd_partials .- sym_partials + err = maximum(abs, diff_array) + + @test err < 1e-6 # or smaller threshold if you want more precision + + fd_bench = @benchmark forwarddiff_partials($η) + sym_bench = @benchmark partial_metric(NormalGamma, NaturalParametersSpace(), $η) + @test min(sym_bench.times...) < min(fd_bench.times...) + #TODO WHY? + # probably smt just wrong here + @test_broken sym_bench.allocs < fd_bench.allocs + end +end + +@testitem "Generic properties of NormalGammaNaturalManifold" begin + import ExponentialFamilyManifolds: NormalGammaNaturalManifold + import ManifoldsBase: manifold_dimension, representation_size, check_point, check_vector, is_point, is_vector + import ManifoldsBase: zero_vector, inner + import ExponentialFamilyManifolds: isproper, NaturalParametersSpace + + using ManifoldsBase, Test + using LinearAlgebra + using ExponentialFamily + + M = NormalGammaNaturalManifold() + valid_point = [1.0, -2.0, 0.5, -1.0] + ef = ExponentialFamilyDistribution(NormalGamma, [1.0, -2.0, 0.5, -1.0], nothing) + + @test @inferred(manifold_dimension(M)) === 4 + @test @inferred(representation_size(M)) === (4,) + @test injectivity_radius(M) === Inf + + @test check_point(M, valid_point) === nothing + @test check_point(M, [2.0, 0.0, 1.0, -3.0]) isa DomainError # invalid because η2>=0 + @test check_point(M, [0.5, -3.0, -1.0, -1.0]) isa DomainError # invalid because η3 <= -1/2 + @test check_point(M, [1, 1, 1]) isa DomainError # invalid because because length(η) != 4 + @test check_point(M, [1, 1, 1, 1, 1]) isa DomainError # invalid because because length(η) != 4 + + # Now define some tangent vectors and test them: + tangent_vectors = [ + [0.0, 0.0, 0.0, 0.0], + [0.1, -0.2, 0.3, -0.4], + [1.0, 1.0, 1.0, 1.0] + ] + + for X in tangent_vectors + @test check_vector(M, valid_point, X) === nothing + end + + @test check_vector(M, valid_point, [1.0, 2.0, 3.0]) isa DomainError + @test check_vector(M, valid_point, [1.0, 2.0, 3.0, 4.0, 5.0]) isa DomainError + + z = zero_vector(M, valid_point) + @test length(z) == 4 + @test all(x->x==0.0, z) + + X = [0.1, 0.0, -0.1, 0.05] + Y = [0.2, -0.3, 0.0, -0.05] + @test inner(M, valid_point, X, Y) ≈ dot(X, fisherinformation(ef), Y) + + @test @eval(@allocated(manifold_dimension($M))) === 0 + @test @eval(@allocated(representation_size($M))) === 0 +end + +@testitem "Simple manifold optimization problem #1" begin + using Manopt, ForwardDiff, Static, StableRNGs, LinearAlgebra, Test + using ExponentialFamily, ExponentialFamilyManifolds + using BayesBase + import ExponentialFamilyManifolds: NormalGammaNaturalManifold + import ManifoldsBase: inner + + rng = StableRNG(42) + p_true = (2.0, 1.5, 1.0, 3.0) + + dist = NormalGamma(p_true...) + ef = convert(ExponentialFamilyDistribution, dist) + η = getnaturalparameters(ef) + data = rand(rng, dist, 200) + + M = NormalGammaNaturalManifold() + + function f(M, p) + ef = ExponentialFamilyDistribution(NormalGamma, p, nothing) + return -mean((x) -> logpdf(ef, x), data) + end + + function grad_f(M, p) + ef = ExponentialFamilyDistribution(NormalGamma, p, nothing) + egrad = ForwardDiff.gradient((p) -> f(M, p), p) + G = fisherinformation(ef) + r = G \ egrad + return r + end + + p0 = [1.0, -2.0, 0.5, -1.0] + + p_opt = gradient_descent( + M, + f, + grad_f, + p0; + stepsize = ConstantLength(0.01), + stopping_criterion = StopWhenGradientNormLess(1e-6), + max_iterations = 500 + ) + + cost_before = f(M, p0) + cost_after = f(M, p_opt) + @test cost_after < cost_before + final_grad = grad_f(M, p_opt) + @test norm(final_grad) < 1e-5 +end \ No newline at end of file From d20a64e73908f1691a699f03842d733706c9551d Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 22 Jan 2025 19:03:43 +0100 Subject: [PATCH 02/12] fix: implement normal gamma natural manifold --- src/ExponentialFamilyManifolds.jl | 1 + src/natural_manifolds/normal_gamma.jl | 15 +++++++++++++++ test/natural_manifolds/normal_gamma_tests.jl | 7 +++++++ 3 files changed, 23 insertions(+) create mode 100644 src/natural_manifolds/normal_gamma.jl create mode 100644 test/natural_manifolds/normal_gamma_tests.jl diff --git a/src/ExponentialFamilyManifolds.jl b/src/ExponentialFamilyManifolds.jl index 1dd76b6..4844b9a 100644 --- a/src/ExponentialFamilyManifolds.jl +++ b/src/ExponentialFamilyManifolds.jl @@ -30,4 +30,5 @@ include("natural_manifolds/wishart.jl") # normal gamma include("normal_gamma/partial.jl") include("normal_gamma/natural.jl") +include("natural_manifolds/normal_gamma.jl") end diff --git a/src/natural_manifolds/normal_gamma.jl b/src/natural_manifolds/normal_gamma.jl new file mode 100644 index 0000000..769f808 --- /dev/null +++ b/src/natural_manifolds/normal_gamma.jl @@ -0,0 +1,15 @@ +""" + get_natural_manifold_base(::Type{NormalGamma}, ::Tuple{}, conditioner=nothing) +Get the natural manifold base for the `NormalGamma` distribution. +""" +function get_natural_manifold_base(::Type{NormalGamma}, ::Tuple{Int}, conditioner=nothing) + return NormalGammaNaturalManifold() +end + +""" + partition_point(::Type{NormalGamma}, ::Tuple{}, p, conditioner=nothing) +Converts the `point` to a compatible representation for the natural manifold of type `NormalGamma`. +""" +function partition_point(::Type{NormalGamma}, ::Tuple{Int}, p, conditioner=nothing) + return p +end diff --git a/test/natural_manifolds/normal_gamma_tests.jl b/test/natural_manifolds/normal_gamma_tests.jl new file mode 100644 index 0000000..7ca30d6 --- /dev/null +++ b/test/natural_manifolds/normal_gamma_tests.jl @@ -0,0 +1,7 @@ +@testitem "Check `NormalGamma` natural manifold" begin + include("natural_manifolds_setuptests.jl") + + test_natural_manifold() do rng + return NormalGamma(10rand(rng), 10rand(rng), 10rand(rng), 10rand(rng)) + end +end From 786593efa902071430dffc5a0efc205d0b04efe8 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 22 Jan 2025 19:06:14 +0100 Subject: [PATCH 03/12] fix: remove NLsolve --- Project.toml | 2 -- 1 file changed, 2 deletions(-) diff --git a/Project.toml b/Project.toml index 16f2892..8416208 100644 --- a/Project.toml +++ b/Project.toml @@ -10,7 +10,6 @@ FastCholesky = "2d5283b6-8564-42b6-bb00-83ed8e915756" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e" ManifoldsBase = "3362f125-f0bb-47a3-aa74-596ffd7ef2fb" -NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" @@ -24,7 +23,6 @@ FastCholesky = "1.3.1" LinearAlgebra = "1.10" Manifolds = "0.9" ManifoldsBase = "0.15" -NLsolve = "4.5.1" Random = "1.10" RecursiveArrayTools = "3" SpecialFunctions = "2.5.0" From 9a89941849594f0b7e8e2bc771817ede014543a9 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 22 Jan 2025 19:16:34 +0100 Subject: [PATCH 04/12] fix: add benchmark tools into tests --- Project.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 8416208..2d70950 100644 --- a/Project.toml +++ b/Project.toml @@ -32,6 +32,7 @@ julia = "1.10" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" @@ -42,4 +43,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Aqua", "Distributions", "JET", "Test", "ReTestItems", "StableRNGs", "StaticArrays", "Manopt", "ForwardDiff"] +test = ["Aqua", "BenchmarkTools", "Distributions", "JET", "Test", "ReTestItems", "StableRNGs", "StaticArrays", "Manopt", "ForwardDiff"] From 1519f729d344dba84cd6aa1d5fcc2fc18c2a2ff0 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 22 Jan 2025 19:50:43 +0100 Subject: [PATCH 05/12] fix: add rand operation for normal gamma manifold --- src/normal_gamma/natural.jl | 30 ++++++++++++++++++++++++++++++ test/normal_gamma_tests.jl | 5 +++++ 2 files changed, 35 insertions(+) diff --git a/src/normal_gamma/natural.jl b/src/normal_gamma/natural.jl index c86c71c..94ce348 100644 --- a/src/normal_gamma/natural.jl +++ b/src/normal_gamma/natural.jl @@ -110,4 +110,34 @@ function inner(::NormalGammaNaturalManifold, p, X, Y) ef = ExponentialFamilyDistribution(NormalGamma, p, nothing) G = fisherinformation(ef) return dot(X, G, Y) +end + +function Random.rand(M::NormalGammaNaturalManifold; kwargs...) + return rand(Random.default_rng(), M; kwargs...) +end + +function Random.rand(rng::AbstractRNG, M::NormalGammaNaturalManifold; kwargs...) + μ = rand(Euclidean())[1] + λ = rand(PositiveNumbers()) + α = rand(PositiveNumbers()) + β = rand(PositiveNumbers()) + return collect(MeanToNatural(NormalGamma)((μ, λ, α, β))) +end + +function Random.rand!(rng::AbstractRNG, M::NormalGammaNaturalManifold, η; kwargs...) + μ = rand(rng, Euclidean())[1] + λ = rand(rng, PositiveNumbers()) + α = rand(rng, PositiveNumbers()) + β = rand(rng, PositiveNumbers()) + + η_nat = MeanToNatural(NormalGamma)((μ, λ, α, β)) + η .= η_nat + return η +end + +import ManifoldsBase: project! + +function ManifoldsBase.project!(M::NormalGammaNaturalManifold, Y, p, X) + Y .= X + return Y end \ No newline at end of file diff --git a/test/normal_gamma_tests.jl b/test/normal_gamma_tests.jl index aa211fe..82019a1 100644 --- a/test/normal_gamma_tests.jl +++ b/test/normal_gamma_tests.jl @@ -104,6 +104,11 @@ end @test @eval(@allocated(manifold_dimension($M))) === 0 @test @eval(@allocated(representation_size($M))) === 0 + + @test check_point(M, rand(M)) === nothing + px = zeros(4) + rand!(M, px) + @test check_point(M, px) === nothing end @testitem "Simple manifold optimization problem #1" begin From 9e80b8639fe45e17757318ccbc2b7d2f7f9753fa Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 22 Jan 2025 20:32:30 +0100 Subject: [PATCH 06/12] fix: use manifolds test set --- Project.toml | 2 ++ src/normal_gamma/natural.jl | 23 +++++++++++++---------- test/normal_gamma_tests.jl | 34 +++++++++++++++++++++++++++++++++- 3 files changed, 48 insertions(+), 11 deletions(-) diff --git a/Project.toml b/Project.toml index 2d70950..84c40b1 100644 --- a/Project.toml +++ b/Project.toml @@ -5,6 +5,7 @@ version = "1.5.0" [deps] BayesBase = "b4ee3484-f114-42fe-b91c-797d54a0c67e" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ExponentialFamily = "62312e5e-252a-4322-ace9-a5f4bf9b357b" FastCholesky = "2d5283b6-8564-42b6-bb00-83ed8e915756" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -18,6 +19,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [compat] BayesBase = "1.3" +Distributions = "0.25.117" ExponentialFamily = "1.6.0" FastCholesky = "1.3.1" LinearAlgebra = "1.10" diff --git a/src/normal_gamma/natural.jl b/src/normal_gamma/natural.jl index 94ce348..8cd0d51 100644 --- a/src/normal_gamma/natural.jl +++ b/src/normal_gamma/natural.jl @@ -1,5 +1,6 @@ using ExponentialFamily using FastCholesky +using Distributions using LinearAlgebra import ManifoldsBase: check_point, check_vector, manifold_dimension, representation_size, exp!, is_point, is_vector @@ -117,19 +118,21 @@ function Random.rand(M::NormalGammaNaturalManifold; kwargs...) end function Random.rand(rng::AbstractRNG, M::NormalGammaNaturalManifold; kwargs...) - μ = rand(Euclidean())[1] - λ = rand(PositiveNumbers()) - α = rand(PositiveNumbers()) - β = rand(PositiveNumbers()) - return collect(MeanToNatural(NormalGamma)((μ, λ, α, β))) + # e.g. draw (μ,λ,α,β) from some easy distributions: + μ = rand(rng, Normal(0,1)) + λ = rand(rng, Exponential(1)) + 0.001 # ensure positivity + α = rand(rng, Exponential(1)) + 0.001 # ensure positivity + β = rand(rng, Exponential(1)) + 0.001 # ensure positivity + # Then transform to natural coords: + η = MeanToNatural(NormalGamma)((μ, λ, α, β)) + return collect(η) end function Random.rand!(rng::AbstractRNG, M::NormalGammaNaturalManifold, η; kwargs...) - μ = rand(rng, Euclidean())[1] - λ = rand(rng, PositiveNumbers()) - α = rand(rng, PositiveNumbers()) - β = rand(rng, PositiveNumbers()) - + μ = rand(rng, Normal(0,1)) + λ = rand(rng, Exponential(1)) + 0.001 # ensure positivity + α = rand(rng, Exponential(1)) + 0.001 # ensure positivity + β = rand(rng, Exponential(1)) + 0.001 # ensure positivity η_nat = MeanToNatural(NormalGamma)((μ, λ, α, β)) η .= η_nat return η diff --git a/test/normal_gamma_tests.jl b/test/normal_gamma_tests.jl index 82019a1..82e72ad 100644 --- a/test/normal_gamma_tests.jl +++ b/test/normal_gamma_tests.jl @@ -158,4 +158,36 @@ end @test cost_after < cost_before final_grad = grad_f(M, p_opt) @test norm(final_grad) < 1e-5 -end \ No newline at end of file +end + +@testitem "Manifolds.test_manifold" begin + using Manifolds, Static, Random + using StableRNGs + + import ExponentialFamilyManifolds: NormalGammaNaturalManifold + + M = NormalGammaNaturalManifold() + rng = StableRNG(42) + + ptss = [ + [ + rand(rng, M), + rand(rng, M), + rand(rng, M), + ] + ] + + for pts in ptss + Manifolds.test_manifold( + M, + pts; + test_exp_log=false, + inverse_retraction_methods=[], # no inverse retraction + retraction_methods=[ExponentialRetraction()], + test_rand_point=true, + test_rand_tvector=false, + parallel_transport=false, + default_inverse_retraction_method=nothing, + ) + end +end From 13d0d1c9f3a2f1f9fc2c2ca74dcec8e6ead94501 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 7 Feb 2025 10:24:33 +0100 Subject: [PATCH 07/12] fix: ensure positivity should quite big, some time drops NaN --- src/normal_gamma/natural.jl | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/src/normal_gamma/natural.jl b/src/normal_gamma/natural.jl index 8cd0d51..e670f89 100644 --- a/src/normal_gamma/natural.jl +++ b/src/normal_gamma/natural.jl @@ -61,7 +61,9 @@ function exp_secondorder( 4D manifold for NormalGamma in natural params η = (η1, η2, η3, η4). Domain constraints, etc. """ -struct NormalGammaNaturalManifold <: AbstractManifold{ℝ} end +struct NormalGammaNaturalManifold <: AbstractManifold{ℝ} + ensure_positivity_shift::Real = 0.1 +end function manifold_dimension(::NormalGammaNaturalManifold) return 4 @@ -120,9 +122,9 @@ end function Random.rand(rng::AbstractRNG, M::NormalGammaNaturalManifold; kwargs...) # e.g. draw (μ,λ,α,β) from some easy distributions: μ = rand(rng, Normal(0,1)) - λ = rand(rng, Exponential(1)) + 0.001 # ensure positivity - α = rand(rng, Exponential(1)) + 0.001 # ensure positivity - β = rand(rng, Exponential(1)) + 0.001 # ensure positivity + λ = rand(rng, Exponential(1)) + M.ensure_positivity_shift + α = rand(rng, Exponential(1)) + M.ensure_positivity_shift + β = rand(rng, Exponential(1)) + M.ensure_positivity_shift # Then transform to natural coords: η = MeanToNatural(NormalGamma)((μ, λ, α, β)) return collect(η) @@ -130,9 +132,9 @@ end function Random.rand!(rng::AbstractRNG, M::NormalGammaNaturalManifold, η; kwargs...) μ = rand(rng, Normal(0,1)) - λ = rand(rng, Exponential(1)) + 0.001 # ensure positivity - α = rand(rng, Exponential(1)) + 0.001 # ensure positivity - β = rand(rng, Exponential(1)) + 0.001 # ensure positivity + λ = rand(rng, Exponential(1)) + M.ensure_positivity_shift + α = rand(rng, Exponential(1)) + M.ensure_positivity_shift + β = rand(rng, Exponential(1)) + M.ensure_positivity_shift η_nat = MeanToNatural(NormalGamma)((μ, λ, α, β)) η .= η_nat return η From 45c87e7ec2637761e6a267263081bc1e494790d5 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 7 Feb 2025 11:19:58 +0100 Subject: [PATCH 08/12] refactor: re-use code from Manifolds.jl --- src/normal_gamma/natural.jl | 86 ++++++++++++++++++------------------- 1 file changed, 43 insertions(+), 43 deletions(-) diff --git a/src/normal_gamma/natural.jl b/src/normal_gamma/natural.jl index e670f89..7d35b0f 100644 --- a/src/normal_gamma/natural.jl +++ b/src/normal_gamma/natural.jl @@ -7,53 +7,29 @@ import ManifoldsBase: check_point, check_vector, manifold_dimension, representat import ManifoldsBase: zero_vector, zero_vector!, default_retraction_method import ManifoldsBase: ExponentialRetraction import ManifoldsBase: injectivity_radius - -""" -Compute all Christoffel symbols Γ^k_{i j}(x). -Returns a 4×4×4 array: Gamma[k,i,j]. -""" -function christoffel(::Type{NormalGamma}, ::NaturalParametersSpace, η) - ef = ExponentialFamilyDistribution(NormalGamma, η, nothing) - G = fisherinformation(ef) - G⁻¹ = cholinv(G) - dG = partial_metric(NormalGamma, NaturalParametersSpace(), η) - # formula: Γ^k_{i,j} = 1/2 * sum_{ℓ} g^{kℓ} ( ∂i g_{ℓj} + ∂j g_{ℓi} - ∂ℓ g_{i j} ) - Γ = zeros(4,4,4) - for k in 1:4 - for i in 1:4 - for j in 1:4 - s = 0.0 - for ℓ in 1:4 - s += G⁻¹[k,ℓ] * (dG[ℓ,j,i] + dG[ℓ,i,j] - dG[i,j,ℓ]) - end - Γ[k,i,j] = 0.5 * s - end - end - end - return Γ -end +import ManifoldsBase: AbstractBasis +import ManifoldsBase: VectorSpaceType, TangentSpaceType +import Manifolds: christoffel_symbols_second function exp_secondorder( ::Type{NormalGamma}, ::NaturalParametersSpace, + Γ, p0, v0 - ) - Γ = christoffel(NormalGamma, NaturalParametersSpace(), p0) - - Δ = zeros(4) - for k in 1:4 - for i in 1:4 - for j in 1:4 - Δ[k] += Γ[k,i,j] * v0[i] * v0[j] - end - end - Δ[k] *= -0.5 - end - +) + Δ = similar(p0) # Preallocate Δ with same type/size as p0 + Manifolds.@einsum Δ[k] = -0.5 * Γ[k,i,j] * v0[i] * v0[j] return p0 + v0 + Δ - end +end + +struct ScoreBasis{ℝ,VST<:VectorSpaceType} <: AbstractBasis{ℝ,VST} + vector_space::VST +end +function ScoreBasis(vst::VST) where {VST<:VectorSpaceType} + return ScoreBasis{ℝ,VST}(vst) +end """ NormalGammaNaturalManifold <: AbstractManifold{ℝ} @@ -61,8 +37,31 @@ function exp_secondorder( 4D manifold for NormalGamma in natural params η = (η1, η2, η3, η4). Domain constraints, etc. """ -struct NormalGammaNaturalManifold <: AbstractManifold{ℝ} - ensure_positivity_shift::Real = 0.1 +struct NormalGammaNaturalManifold{ℝ, T, B} <: AbstractManifold{ℝ} + ensure_positivity_shift::T + basis::B +end + +function NormalGammaNaturalManifold() + return NormalGammaNaturalManifold{ℝ, Float64, ScoreBasis{ℝ, TangentSpaceType}}(0.1, ScoreBasis(TangentSpaceType())) +end + +function Manifolds.local_metric( + ::NormalGammaNaturalManifold, + p, + ::ScoreBasis{ℝ, TangentSpaceType} +) + ef = ExponentialFamilyDistribution(NormalGamma, p, nothing) + return fisherinformation(ef) +end + +function Manifolds.local_metric_jacobian( + M::NormalGammaNaturalManifold, + p, + ::ScoreBasis{ℝ, TangentSpaceType}; + kwargs... +) + return partial_metric(NormalGamma, NaturalParametersSpace(), p) end function manifold_dimension(::NormalGammaNaturalManifold) @@ -88,8 +87,9 @@ function check_vector(::NormalGammaNaturalManifold, η, X) return nothing end -function exp!(::NormalGammaNaturalManifold, η_out::AbstractVector, η_in::AbstractVector, X_in::AbstractVector, t::Real=1.0) - η_out .= exp_secondorder(NormalGamma, NaturalParametersSpace(), η_in, t.*X_in) +function exp!(M::NormalGammaNaturalManifold, η_out::AbstractVector, η_in::AbstractVector, X_in::AbstractVector, t::Real=1.0) + Γ = christoffel_symbols_second(M, η_in, M.basis) + η_out .= exp_secondorder(NormalGamma, NaturalParametersSpace(), Γ, η_in, t.*X_in) return η_out end From 45c13930fa11003055be49ad384ed08c64ed8335 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 7 Feb 2025 11:31:20 +0100 Subject: [PATCH 09/12] fix: normal gamma is an open set in R^{4}, check that projection is indentity --- test/normal_gamma_tests.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/test/normal_gamma_tests.jl b/test/normal_gamma_tests.jl index 82e72ad..72a39e3 100644 --- a/test/normal_gamma_tests.jl +++ b/test/normal_gamma_tests.jl @@ -190,4 +190,8 @@ end default_inverse_retraction_method=nothing, ) end + + point = rand(M) + vector = rand(rng, 4) + @test isapprox(vector, project(M, point, vector)) end From 1cbf2a0a1870fb822ec4ed5504adb6e2676801da Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 7 Feb 2025 11:45:01 +0100 Subject: [PATCH 10/12] docs: add NormalGammaNaturalManifold into docs --- docs/src/index.md | 1 + src/normal_gamma/natural.jl | 7 ++++--- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/src/index.md b/docs/src/index.md index 1002d99..8a03ecc 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -82,6 +82,7 @@ ExponentialFamilyManifolds.ShiftedPositiveNumbers ExponentialFamilyManifolds.ShiftedNegativeNumbers ExponentialFamilyManifolds.SymmetricNegativeDefinite ExponentialFamilyManifolds.SinglePointManifold +ExponentialFamilyManifolds.NormalGammaNaturalManifold ``` ## Optimization example diff --git a/src/normal_gamma/natural.jl b/src/normal_gamma/natural.jl index 7d35b0f..e9e736b 100644 --- a/src/normal_gamma/natural.jl +++ b/src/normal_gamma/natural.jl @@ -34,8 +34,9 @@ end """ NormalGammaNaturalManifold <: AbstractManifold{ℝ} -4D manifold for NormalGamma in natural params η = (η1, η2, η3, η4). -Domain constraints, etc. +4D manifold for NormalGamma in natural params η = (η1, η2, η3, η4) with fisher information metric. + +(see ExponentialFamily.jl/src/normal_gamma.jl) """ struct NormalGammaNaturalManifold{ℝ, T, B} <: AbstractManifold{ℝ} ensure_positivity_shift::T @@ -142,7 +143,7 @@ end import ManifoldsBase: project! -function ManifoldsBase.project!(M::NormalGammaNaturalManifold, Y, p, X) +function ManifoldsBase.project!(::NormalGammaNaturalManifold, Y, p, X) Y .= X return Y end \ No newline at end of file From 86da7fdb2d28b47ae49b00d1397938cc70db9157 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 7 Feb 2025 11:56:53 +0100 Subject: [PATCH 11/12] docs: improve NormalGammaNaturalManifold docstring --- src/normal_gamma/natural.jl | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/src/normal_gamma/natural.jl b/src/normal_gamma/natural.jl index e9e736b..e83c5ab 100644 --- a/src/normal_gamma/natural.jl +++ b/src/normal_gamma/natural.jl @@ -32,11 +32,21 @@ function ScoreBasis(vst::VST) where {VST<:VectorSpaceType} end """ - NormalGammaNaturalManifold <: AbstractManifold{ℝ} + NormalGammaNaturalManifold{ℝ, T, B} <: AbstractManifold{ℝ} -4D manifold for NormalGamma in natural params η = (η1, η2, η3, η4) with fisher information metric. +A 4-dimensional manifold representing the natural parameter space of the Normal-Gamma distribution. -(see ExponentialFamily.jl/src/normal_gamma.jl) +The manifold is parameterized by natural parameters η = (η₁, η₂, η₃, η₄) and equipped with the Fisher information metric. +The metric is computed in the score basis, which is defined as derivatives of the score vector with respect to the +natural parameters (see [An elementary introduction to information geometry](https://arxiv.org/pdf/1808.08271), equation 79). + +To ensure valid sampling, the natural parameters are shifted by a small positive number to maintain positivity constraints. + +# Fields +- `ensure_positivity_shift::T`: Small positive number used to shift parameters during sampling +- `basis::B`: Score basis used for computing the Fisher information metric + +See also [`ExponentialFamily.NormalGamma`](@ref) for details about the Normal-Gamma distribution implementation. """ struct NormalGammaNaturalManifold{ℝ, T, B} <: AbstractManifold{ℝ} ensure_positivity_shift::T From f9beac2787404ab960c70030fb4fe7d46e75836a Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 7 Feb 2025 12:48:28 +0100 Subject: [PATCH 12/12] docs: point to correct place for NormalGamma distribution --- src/normal_gamma/natural.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/normal_gamma/natural.jl b/src/normal_gamma/natural.jl index e83c5ab..2021623 100644 --- a/src/normal_gamma/natural.jl +++ b/src/normal_gamma/natural.jl @@ -46,7 +46,7 @@ To ensure valid sampling, the natural parameters are shifted by a small positive - `ensure_positivity_shift::T`: Small positive number used to shift parameters during sampling - `basis::B`: Score basis used for computing the Fisher information metric -See also [`ExponentialFamily.NormalGamma`](@ref) for details about the Normal-Gamma distribution implementation. +See also [`ExponentialFamily.NormalGamma`](https://reactivebayes.github.io/ExponentialFamily.jl/stable/library/#ExponentialFamily.NormalGamma) for details about the Normal-Gamma distribution implementation. """ struct NormalGammaNaturalManifold{ℝ, T, B} <: AbstractManifold{ℝ} ensure_positivity_shift::T