Add natural gamma manifold

Project.toml

@@ -5,27 +5,36 @@ 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"
 Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
 ManifoldsBase = "3362f125-f0bb-47a3-aa74-596ffd7ef2fb"
 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"
+Distributions = "0.25.117"
 ExponentialFamily = "1.6.0"
+FastCholesky = "1.3.1"
 LinearAlgebra = "1.10"
 Manifolds = "0.9"
 ManifoldsBase = "0.15"
 Random = "1.10"
 RecursiveArrayTools = "3"
+SpecialFunctions = "2.5.0"
 Static = "0.8, 1"
+StaticArrays = "1.9.10"
 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"
@@ -36,4 +45,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"]

docs/src/index.md

@@ -82,6 +82,7 @@ ExponentialFamilyManifolds.ShiftedPositiveNumbers
 ExponentialFamilyManifolds.ShiftedNegativeNumbers
 ExponentialFamilyManifolds.SymmetricNegativeDefinite
 ExponentialFamilyManifolds.SinglePointManifold
+ExponentialFamilyManifolds.NormalGammaNaturalManifold
 ```
 
 ## Optimization example

src/ExponentialFamilyManifolds.jl

@@ -27,4 +27,8 @@ 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")
+include("natural_manifolds/normal_gamma.jl")
 end

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

src/normal_gamma/natural.jl

@@ -0,0 +1,159 @@
+using ExponentialFamily
+using FastCholesky
+using Distributions
+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
+import ManifoldsBase: AbstractBasis
+import ManifoldsBase: VectorSpaceType, TangentSpaceType
+import Manifolds: christoffel_symbols_second
+
+function exp_secondorder(
+    ::Type{NormalGamma},
+    ::NaturalParametersSpace,
+    Γ,
+    p0,
+    v0
+)   
+    Δ = 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
+
+struct ScoreBasis{ℝ,VST<:VectorSpaceType} <: AbstractBasis{ℝ,VST}
+    vector_space::VST
+end
+
+function ScoreBasis(vst::VST) where {VST<:VectorSpaceType}
+    return ScoreBasis{ℝ,VST}(vst)
+end
+
+"""
+    NormalGammaNaturalManifold{ℝ, T, B} <: AbstractManifold{ℝ}
+
+A 4-dimensional manifold representing the natural parameter space of the Normal-Gamma distribution.
+
+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`](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
+    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)
+    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!(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
+
+# 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
+
+function Random.rand(M::NormalGammaNaturalManifold; kwargs...)
+    return rand(Random.default_rng(), M; kwargs...)
+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)) + 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(η)
+end
+
+function Random.rand!(rng::AbstractRNG, M::NormalGammaNaturalManifold, η; kwargs...)
+    μ = rand(rng, Normal(0,1))
+    λ = 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 η
+end
+
+import ManifoldsBase: project!
+
+function ManifoldsBase.project!(::NormalGammaNaturalManifold, Y, p, X)
+    Y .= X
+    return Y
+end

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