diff --git a/.gitignore b/.gitignore index 92b50fe..9082aaf 100644 --- a/.gitignore +++ b/.gitignore @@ -4,4 +4,4 @@ /Manifest.toml /docs/Manifest.toml /docs/build/ -*Manifest.toml +scripts/Manifest.toml diff --git a/Project.toml b/Project.toml index 7fa596d..f144c51 100644 --- a/Project.toml +++ b/Project.toml @@ -12,6 +12,7 @@ ManifoldsBase = "3362f125-f0bb-47a3-aa74-596ffd7ef2fb" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [compat] ADTypes = "1.14.0" @@ -27,6 +28,7 @@ ManifoldsBase = "2.2.0" Random = "1.10" RecursiveArrayTools = "3" Static = "0.8, 1" +StaticArrays = "1.9.15" julia = "1.10" [extras] @@ -34,8 +36,8 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" -ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" ManifoldDiff = "af67fdf4-a580-4b9f-bbec-742ef357defd" diff --git a/src/ExponentialFamilyManifolds.jl b/src/ExponentialFamilyManifolds.jl index 51b6fd7..7aa1403 100644 --- a/src/ExponentialFamilyManifolds.jl +++ b/src/ExponentialFamilyManifolds.jl @@ -1,9 +1,11 @@ module ExponentialFamilyManifolds -using BayesBase, ExponentialFamily, ManifoldsBase, Manifolds, Random, LinearAlgebra +using BayesBase, + ExponentialFamily, ManifoldsBase, Manifolds, Random, LinearAlgebra, StaticArrays include("single_point_manifold.jl") include("natural_manifolds.jl") +include("with_metric.jl") include("natural_manifolds/bernoulli.jl") include("natural_manifolds/beta.jl") @@ -24,4 +26,9 @@ include("natural_manifolds/poisson.jl") include("natural_manifolds/rayleigh.jl") include("natural_manifolds/weibull.jl") include("natural_manifolds/wishart.jl") + +# metric Manifolds +include("metric_manifolds/jacobians.jl") +include("metric_manifolds/bernoulli.jl") +include("metric_manifolds/normal.jl") end diff --git a/src/metric_manifolds/bernoulli.jl b/src/metric_manifolds/bernoulli.jl new file mode 100644 index 0000000..f11ae3e --- /dev/null +++ b/src/metric_manifolds/bernoulli.jl @@ -0,0 +1,161 @@ +using ManifoldsBase + +# Stop forwarding for functions we implement +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.exp), ::Type +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.exp!), ::Type +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.exp_fused), ::Type +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.exp_fused!), ::Type +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.log), ::Type +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.log!), ::Type +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_to), + ::Type, +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_to!), + ::Type, +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_direction), + ::Type, +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_direction!), + ::Type, +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, + ::typeof(ManifoldsBase.vector_transport_to), + ::Type, +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, + ::typeof(ManifoldsBase.vector_transport_to!), + ::Type, +) where {F} = _STOP + +# Also catch the single-argument versions (without Type) +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.exp) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.exp!) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.exp_fused) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.exp_fused!) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.log) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.log!) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.parallel_transport_to) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.parallel_transport_to!) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_direction), +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_direction!), +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.vector_transport_to) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,Bernoulli,NaturalMetric}, ::typeof(ManifoldsBase.vector_transport_to!) +) where {F} = _STOP + +function _geodesic_exact_bernoulli(t, x0, v0) + # Compute metric at initial point + g0 = exp(x0) / (1 + exp(x0))^2 + + # Energy constant + E = g0 * v0^2 + + # Integration constant from initial condition + C = 2 * atan(exp(x0/2)) + + # Sign based on initial velocity + sign_v = sign(v0) + + # Position + x = 2 * log(abs(tan((sign_v * sqrt(E) * t + C) / 2))) + + # Velocity from energy conservation + g_x = exp(x) / (1 + exp(x))^2 + v = sign_v * sqrt(E / g_x) + + return x, v +end + +function _log_map_bernoulli(x0, x1) + C = 2 * atan(exp(x0/2)) + C1 = 2 * atan(exp(x1/2)) + + # Determine sign (direction) + sign_v = sign(C1 - C) + + # Solve for sqrt(E) + sqrt_E = abs(C1 - C) + + # Compute v0 from E = g(x0) * v0^2 + g0 = exp(x0) / (1 + exp(x0))^2 + v0 = sign_v * sqrt(sqrt_E^2 / g0) + + return v0 +end + +function ManifoldsBase.exp_fused!( + ::WithMetric{F,Bernoulli,NaturalMetric}, q, p, X, t::Number +) where {F} + x, _ = _geodesic_exact_bernoulli(t, p[1], X[1]) + q .= x + return q +end + +function ManifoldsBase.log!(::WithMetric{F,Bernoulli,NaturalMetric}, X, p, q) where {F} + X .= _log_map_bernoulli(p[1], q[1]) + return X +end + +_sqrtg_eta(η) = 0.5 .* sech.(η ./ 2) + +function ManifoldsBase.parallel_transport_to( + ::WithMetric{F,Bernoulli,NaturalMetric}, p, X, q; kwargs... +) where {F} + s_p = _sqrtg_eta(p) + s_q = _sqrtg_eta(q) + return X .* (s_p ./ s_q) +end + +function ManifoldsBase.parallel_transport_to!( + M::WithMetric{F,Bernoulli,NaturalMetric}, Y, p, X, q; kwargs... +) where {F} + Y .= ManifoldsBase.parallel_transport_to(M, p, X, q; kwargs...) + return Y +end diff --git a/src/metric_manifolds/jacobians.jl b/src/metric_manifolds/jacobians.jl new file mode 100644 index 0000000..10b3572 --- /dev/null +++ b/src/metric_manifolds/jacobians.jl @@ -0,0 +1,93 @@ +""" + jacobian_manifold_to_nat!(M, X_nat, X_manifold) + +Transform tangent vector from manifold coordinates to natural parameter coordinates. +""" +function jacobian_manifold_to_nat( + M::ExponentialFamilyManifolds.NaturalParametersManifold, X_manifold +) + X_nat = jacobian_manifold_to_nat!(M, zero_vector(M, X_manifold), X_manifold) + return X_nat +end + +""" + jacobian_nat_to_manifold(M, X_nat) + +Transform tangent vector from natural parameter coordinates to manifold coordinates. +""" +function jacobian_nat_to_manifold( + M::ExponentialFamilyManifolds.NaturalParametersManifold, X_nat +) + X_manifold = jacobian_nat_to_manifold!(M, zero_vector(M, X_nat), X_nat) + return X_manifold +end + +""" + jacobian_nat_to_manifold!(M, X_manifold, X_nat) + +Transform tangent vector from natural parameter coordinates to default manifold coordinates. +For NormalMeanVariance: (dη₁, dη₂) → (dη₁, dλ) where dλ = -dη₂. +""" +function jacobian_nat_to_manifold!( + ::ExponentialFamilyManifolds.NaturalParametersManifold{ + F,ExponentialFamily.NormalMeanVariance + }, + X_manifold, + X_nat, +) where {F} + X_manifold[1:1] .= X_nat[1] + X_manifold[2:2] .= -X_nat[2] + return X_manifold +end + +""" + jacobian_manifold_to_nat!(M, X_nat, X_manifold) + +Transform tangent vector from manifold coordinates to natural parameter coordinates. +For NormalMeanVariance: (dη₁, dλ) → (dη₁, dη₂) where dη₂ = -dλ. +""" +function jacobian_manifold_to_nat!( + ::ExponentialFamilyManifolds.NaturalParametersManifold{ + F,ExponentialFamily.NormalMeanVariance + }, + X_nat, + X_manifold, +) where {F} + X_nat[1:1] .= X_manifold[1] + X_nat[2:2] .= -X_manifold[2] + return X_nat +end + +""" + jacobian_nat_to_manifold!(M, X_manifold, X_nat) + +Transform tangent vector from natural parameter coordinates to default manifold coordinates. +For NormalMeanVariance: (dη₁, dη₂) → (dη₁, dλ) where dλ = -dη₂. +""" +function jacobian_nat_to_manifold!( + ::ExponentialFamilyManifolds.NaturalParametersManifold{ + F,ExponentialFamily.Bernoulli + }, + X_manifold, + X_nat, +) where {F} + X_manifold .= X_nat[1] + return X_manifold +end + +""" + jacobian_manifold_to_nat!(M, X_nat, X_manifold) + +Transform tangent vector from manifold coordinates to natural parameter coordinates. +For NormalMeanVariance: (dη₁, dλ) → (dη₁, dη₂) where dη₂ = -dλ. +""" +function jacobian_manifold_to_nat!( + ::ExponentialFamilyManifolds.NaturalParametersManifold{ + F,ExponentialFamily.Bernoulli + }, + X_nat, + X_manifold, +) where {F} + X_nat .= X_manifold + return X_nat +end \ No newline at end of file diff --git a/src/metric_manifolds/normal.jl b/src/metric_manifolds/normal.jl new file mode 100644 index 0000000..d1f65f6 --- /dev/null +++ b/src/metric_manifolds/normal.jl @@ -0,0 +1,419 @@ +using StaticArrays +using RecursiveArrayTools: ArrayPartition + +# For NormalMeanVariance with NaturalMetric we also override geodesic / log / transport +# related functions and curvature in `metric_manifolds/normal.jl`. Mark all of these as +# non-forwarding for that specific case so they are never delegated to the decorated manifold. +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, ::typeof(ManifoldsBase.exp), ::Type +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, ::typeof(ManifoldsBase.exp!), ::Type +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.exp_fused), + ::Type, +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.exp_fused!), + ::Type, +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, ::typeof(ManifoldsBase.log), ::Type +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, ::typeof(ManifoldsBase.log!), ::Type +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_to), + ::Type, +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_to!), + ::Type, +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_direction), + ::Type, +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_direction!), + ::Type, +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.vector_transport_to), + ::Type, +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.vector_transport_to!), + ::Type, +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.sectional_curvature_max), + ::Type, +) where {F} = _STOP + +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, ::typeof(ManifoldsBase.exp) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, ::typeof(ManifoldsBase.exp!) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, ::typeof(ManifoldsBase.exp_fused) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, ::typeof(ManifoldsBase.exp_fused!) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, ::typeof(ManifoldsBase.log) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, ::typeof(ManifoldsBase.log!) +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_to), +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_to!), +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_direction), +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.parallel_transport_direction!), +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.vector_transport_to), +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.vector_transport_to!), +) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + ::typeof(ManifoldsBase.sectional_curvature_max), +) where {F} = _STOP + +# ============================================================================ +# Helper functions: coordinate conversions and hyperbolic geometry +# ============================================================================ +# ------------------------------------------------------------------ +# Representation: +# manifold point p = (η₁, λ), where λ = -η₂ > 0 +# natural parameters η = (η₁, η₂), η₂ < 0 +# ------------------------------------------------------------------ + +@inline _point_to_eta(p) = SA[p[1], -p[2]] # (η₁, λ) -> (η₁, η₂) +@inline _eta_to_point(η::SVector{2}) = SA[η[1], -η[2]] + +@inline _X_to_eta(X) = SA[X[1], -X[2]] # tangent: dη₁ = dη₁, dη₂ = -dλ +@inline _eta_to_X(Xη::SVector{2}) = SA[Xη[1], -Xη[2]] + +@inline function _eta_to_meanvariance(η1, η2) + μ = -η1 / (2 * η2) + σ² = -inv(2 * η2) + return μ, σ² +end + +@inline function _meanvariance_to_eta(μ, σ²) + η1 = μ / σ² + η2 = -inv(2 * σ²) + return η1, η2 +end + +# (η₁, η₂) -> Poincaré half-plane coords (x,y), y>0 +# here x = μ/√2, y = σ +@inline function _eta_to_halfplane(η::SVector{2,T}) where {T} + η1, η2 = η + μ, σ² = _eta_to_meanvariance(η1, η2) + @assert σ² > zero(T) "NormalMeanVariance natural parameters must satisfy η₂ < 0" + σ = sqrt(σ²) + x = μ / sqrt(T(2)) + y = σ + return SA[x, y] +end + +# (x,y) in half-plane -> (η₁, η₂) +@inline function _halfplane_to_eta(xy::SVector{2,T}) where {T} + x, y = xy + μ = sqrt(T(2)) * x + σ² = y^2 + η1, η2 = _meanvariance_to_eta(μ, σ²) + return SA[η1, η2] +end + +# --- Jacobians between natural params and half-plane coords --- + +# J_η→(x,y) evaluated at η = (η₁, η₂) +# [∂x/∂η₁ ∂x/∂η₂] +# [∂y/∂η₁ ∂y/∂η₂] +@inline function _J_eta_to_xy(η::SVector{2}) + η1, η2 = η + μ, σ² = _eta_to_meanvariance(η1, η2) + σ = sqrt(σ²) + c = sqrt(2.0) + a11 = -c / (4 * η2) + a12 = c * η1 / (4 * η2^2) + a21 = 0.0 + a22 = -σ / (2 * η2) # equals sqrt(2)/(4*(-η2)^(3/2)) + return SA[ + a11 a12; + a21 a22 + ] +end + +# J_(x,y)→η evaluated at (x,y) +# [∂η₁/∂x ∂η₁/∂y] +# [∂η₂/∂x ∂η₂/∂y] +@inline function _J_xy_to_eta(xy::SVector{2}) + x, y = xy + c = sqrt(2.0) + y2 = y^2 + y3 = y2 * y + b11 = c / y2 + b12 = -2c * x / y3 + b21 = 0.0 + b22 = 1.0 / y3 + return SA[ + b11 b12; + b21 b22 + ] +end + +# --- Half-plane ↔ hyperboloid (curvature -1) --- + +# (x,y) in upper half-plane -> hyperboloid point u ∈ ℝ³ +# with Minkowski metric ⟨u,u⟩ = -u₀² + u₁² + u₂² = -1, u₀>0 +@inline function _hp_to_hyp(xy::SVector{2}) + x, y = xy + x2 = x^2 + y2 = y^2 + u0 = (x2 + y2 + 1) / (2*y) + u1 = x / y + u2 = (x2 + y2 - 1) / (2*y) + return SA[u0, u1, u2] +end + +# hyperboloid point -> (x,y) in half-plane +@inline function _hyp_to_hp(u::SVector{3}) + u0, u1, u2 = u + denom = u0 - u2 + x = u1 / denom + y = 1 / denom + return SA[x, y] +end + +# J_(x,y)→u (3×2) +@inline function _J_hp_to_hyp(xy::SVector{2}) + x, y = xy + x2 = x^2 + y2 = y^2 + return SA[ + x/y (-x2 + y2 - 1) / (2*y2); + 1/y -x / y2; + x/y (-x2 + y2 + 1) / (2*y2) + ] +end + +# J_u→(x,y) (2×3) = derivative of inverse map +@inline function _J_hyp_to_hp(u::SVector{3}) + u0, u1, u2 = u + denom = u0 - u2 + denom2 = denom^2 + return SA[ + -u1/denom2 1/denom u1/denom2; + -1/denom2 0.0 1/denom2 + ] +end + +# --- Hyperbolic geometry on the hyperboloid --- + +@inline _mdot(a::SVector{3}, b::SVector{3}) = -a[1]*b[1] + a[2]*b[2] + a[3]*b[3] # Minkowski inner product + +# Exponential on the hyperboloid (curvature -1) +function _hyperbolic_exp(u::SVector{3,T}, W::SVector{3,T}, t::T) where {T<:Real} + m2 = _mdot(W, W) + m2 = max(m2, zero(T)) + speed = sqrt(m2) + if speed < eps(T) + return u + end + ct = cosh(speed*t) + st = sinh(speed*t) + u_t = ct*u + (st/speed)*W + return u_t +end + +# Logarithmic map on the hyperboloid (Nagano et al., Eq. (6)) +function _hyperbolic_log(u::SVector{3,T}, v::SVector{3,T}) where {T<:Real} + if u == v + return zero(u) + end + ip = _mdot(u, v) + α = -ip + d = acosh(α) + w = v - α*u # v - α u, α = -⟨u,v⟩ + n2 = _mdot(w, w) + if n2 <= zero(T) + return zero(u) + end + n = sqrt(n2) + return (d/n) * w +end + +# Parallel transport on the hyperboloid (Nagano et al., Eq. (3)) +function _hyperbolic_parallel_transport( + ν::SVector{3,T}, μ::SVector{3,T}, v::SVector{3,T} +) where {T<:Real} + α = -_mdot(ν, μ) + num = _mdot(μ - α*ν, v) + coeff = num / (α + one(T)) + return v + coeff * (ν + μ) +end + +# ============================================================================ +# High-level geodesic and log map operations in natural coordinates +# ============================================================================ + +# Geodesic in η-coordinates (internal helper) +function geodesic_exact_eta(t::Real, η0::SVector{2,T}, Xη0::SVector{2,T}) where {T<:Real} + # point: η -> (x,y) -> hyperboloid + xy0 = _eta_to_halfplane(η0) + u0 = _hp_to_hyp(xy0) + + # tangent: η -> (x,y) -> hyperboloid + # Fisher metric = 2 × Poincaré, so scale by 1/√2 when pushing to hyperboloid + Jη_xy = _J_eta_to_xy(η0) + V0 = Jη_xy * Xη0 + Jxy_u = _J_hp_to_hyp(xy0) + W0 = (Jxy_u * V0) / sqrt(T(2)) + + # geodesic on hyperboloid + u_t = _hyperbolic_exp(u0, W0, T(t)) + + # back to η + xy_t = _hyp_to_hp(u_t) + η_t = _halfplane_to_eta(xy_t) + + return η_t +end + +# Public geodesic: manifold coordinates p = (η₁, λ) +function geodesic_exact(t::Real, p, X) + η0 = _point_to_eta(p) + Xη0 = _X_to_eta(X) + η_t = geodesic_exact_eta(t, η0, Xη0) + p_t = _eta_to_point(η_t) + return p_t +end + +# Log map in η-coordinates +function log_map_eta(ηp::SVector{2,T}, ηq::SVector{2,T}) where {T<:Real} + # points to half-plane and hyperboloid + xyp = _eta_to_halfplane(ηp) + xyq = _eta_to_halfplane(ηq) + up = _hp_to_hyp(xyp) + uq = _hp_to_hyp(xyq) + + # log on hyperboloid + ξ = _hyperbolic_log(up, uq) + + # back to (x,y) then η (tangent at p) + # Fisher metric = 2 × Poincaré, so scale by √2 when pulling from hyperboloid + Ju_xy = _J_hyp_to_hp(up) + Vp = Ju_xy * ξ + Jxy_η = _J_xy_to_eta(xyp) + Xη = (Jxy_η * Vp) * sqrt(T(2)) + + return Xη +end + +# Log map in manifold coordinates p = (η₁, λ) +function log_map(p, q) + ηp = _point_to_eta(p) + ηq = _point_to_eta(q) + Xη = log_map_eta(ηp, ηq) + X = _eta_to_X(Xη) + return X +end + +function ManifoldsBase.exp_fused!( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, q, p, X, t::Number +) where {F} + p_t = geodesic_exact(t, p, X) + q[1] = p_t[1] + q[2] = p_t[2] + return q +end + +function ManifoldsBase.log!( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, X, p, q +) where {F} + X .= log_map(p, q) + return X +end + +function ManifoldsBase.parallel_transport_to( + M::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, q; kwargs... +) where {F} + ηp = _point_to_eta(p) + ηq = _point_to_eta(q) + Xη = _X_to_eta(X) + + # p, q to hyperboloid + xyp = _eta_to_halfplane(ηp) + xyq = _eta_to_halfplane(ηq) + up = _hp_to_hyp(xyp) + uq = _hp_to_hyp(xyq) + + # push X to hyperboloid (scale by 1/√2 for metric) + Jη_xy = _J_eta_to_xy(ηp) + Vp = Jη_xy * Xη + Jxy_u = _J_hp_to_hyp(xyp) + Wp = (Jxy_u * Vp) / sqrt(2.0) + + # parallel transport on hyperboloid + Wq = _hyperbolic_parallel_transport(up, uq, Wp) + + # pull back to η at q (scale by √2 for metric) + Ju_xy = _J_hyp_to_hp(uq) + Vq = Ju_xy * Wq + Jxy_η = _J_xy_to_eta(xyq) + Xηq = (Jxy_η * Vq) * sqrt(2.0) + + Xq = _eta_to_X(Xηq) + return ArrayPartition([Xq[1]], [Xq[2]]) +end + +function ManifoldsBase.parallel_transport_to!( + M::WithMetric{F,NormalMeanVariance,NaturalMetric}, Y, p, X, q; kwargs... +) where {F} + Y .= ManifoldsBase.parallel_transport_to(M, p, X, q; kwargs...) + return Y +end + +""" +Sectional curvature for Normal manifold +The Normal Fisher manifold is isometric to hyperbolic space with curvature -1/2 +(Fisher metric = 2 × Poincaré metric, curvature scales as κ/c for metric c·g) +""" +function ManifoldsBase.sectional_curvature_max( + ::WithMetric{F,NormalMeanVariance,NaturalMetric} +) where {F} + return -0.5 +end diff --git a/src/natural_manifolds.jl b/src/natural_manifolds.jl index 7edd8d4..46525a4 100644 --- a/src/natural_manifolds.jl +++ b/src/natural_manifolds.jl @@ -24,7 +24,7 @@ getconditioner(M::NaturalParametersManifold) = M.conditioner # Forward all functions by default to the decorated base manifold @inline ManifoldsBase.get_forwarding_type(::NaturalParametersManifold, ::Any) = ManifoldsBase.SimpleForwardingType() -@inline ManifoldsBase.get_forwarding_type(::ExponentialFamilyManifolds.NaturalParametersManifold, ::Any, ::Type) = ManifoldsBase.SimpleForwardingType() +@inline ManifoldsBase.get_forwarding_type(::NaturalParametersManifold, ::Any, ::Type) = ManifoldsBase.SimpleForwardingType() function ExponentialFamily.exponential_family_typetag( ::NaturalParametersManifold{𝔽,T} diff --git a/src/with_metric.jl b/src/with_metric.jl new file mode 100644 index 0000000..6c929f3 --- /dev/null +++ b/src/with_metric.jl @@ -0,0 +1,70 @@ +const _STOP = ManifoldsBase.StopForwardingType() +const _FWD = ManifoldsBase.SimpleForwardingType() + +abstract type MetricMode end +struct NaturalMetric <: MetricMode end # Fisher info in active coordinates + +""" + WithMetric(::Type{T}, metric, dims, base, conditioner) + +The manifold for the natural parameters of the distribution of type `T` with dimensions `dims`, equipeed with metric. +An internal structure, use `get_fisher_manifold` to create an instance of a manifold for the natural parameters of distribution of type `T`. + +The key idea of this manifold that for it checked that goedesic is compatible with metric. +For noe the main use is with `NaturalMetric` which is Fisher information in natural coordinates. +""" +struct WithMetric{𝔽,T,Mode<:MetricMode,M<:NaturalParametersManifold} <: + AbstractDecoratorManifold{𝔽} + man::M +end +ManifoldsBase.decorated_manifold(W::WithMetric) = W.man + +function get_fisher_manifold(::Type{T}, dims, conditioner=nothing) where {T} + natural_manifold = ExponentialFamilyManifolds.get_natural_manifold(T, dims, conditioner) + fisher_manifold = with_natural_metric(natural_manifold) + return fisher_manifold +end + +function with_natural_metric(M::NaturalParametersManifold{F}) where {F} + ef_typetag = ExponentialFamily.exponential_family_typetag(M) + WithMetric{F,ef_typetag,NaturalMetric,typeof(M)}(M) +end + +# When NaturalMetric is active, do NOT forward inner/norm. We implement them on the wrapper. +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,T,NaturalMetric}, ::typeof(inner) +) where {F,T} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,T,NaturalMetric}, ::typeof(norm) +) where {F,T} = _STOP + +# Everything else still forwards. +@inline ManifoldsBase.get_forwarding_type(::WithMetric, ::Any) = _FWD +@inline ManifoldsBase.get_forwarding_type(::WithMetric, ::Any, ::Type) = _FWD + +ExponentialFamily.exponential_family_typetag(::WithMetric{F,T}) where {F,T} = T + +function Base.convert(::Type{ExponentialFamilyDistribution}, M::WithMetric, p) + natural_manifold = M.man + return ExponentialFamilyDistribution( + exponential_family_typetag(natural_manifold), + transform_back(natural_manifold, p), + getconditioner(natural_manifold), + nothing, + ) +end + +function ManifoldsBase.norm(M::WithMetric{F,T,NaturalMetric}, p, X) where {F,T} + return sqrt(inner(M, p, X, X)) +end + +function ManifoldsBase.inner( + M::WithMetric{F,T,NaturalMetric}, p, X, Y +) where {F,T} + natural_M = M.man + Xη = jacobian_manifold_to_nat(natural_M, X) + Yη = jacobian_manifold_to_nat(natural_M, Y) + ef = convert(ExponentialFamilyDistribution, M.man, p) + fisher_info = fisherinformation(ef) + return dot(Xη, fisher_info, Yη) +end diff --git a/test/metric_manifolds/bernoulli_tests.jl b/test/metric_manifolds/bernoulli_tests.jl new file mode 100644 index 0000000..3e83aef --- /dev/null +++ b/test/metric_manifolds/bernoulli_tests.jl @@ -0,0 +1,13 @@ +@testitem "Check `Bernoulli` fisher manifold" begin + include("metric_manifolds_setuptests.jl") + test_metric_manifold(maximal_norm=0.3) do rng + return Bernoulli(rand(rng)) + end +end + +@testitem "Check `Bernoulli` fisher manifold MLE works" begin + include("mle_metric_manifolds_setuptests.jl") + test_mle_works() do rng + return Bernoulli(rand(rng)) + end +end diff --git a/test/metric_manifolds/metric_manifolds_setuptests.jl b/test/metric_manifolds/metric_manifolds_setuptests.jl new file mode 100644 index 0000000..1880888 --- /dev/null +++ b/test/metric_manifolds/metric_manifolds_setuptests.jl @@ -0,0 +1,84 @@ +using StableRNGs, ExponentialFamily, ManifoldsBase, LinearAlgebra +using Distributions + +import ExponentialFamilyManifolds: get_fisher_manifold, partition_point + +function test_metric_manifold( + f; seed=42, ndistributions=100, test_points=10, maximal_norm=1, tol=1e-3 +) + rng = StableRNG(seed) + + foreach(1:ndistributions) do _ + distribution = f(rng) + sample = rand(rng, distribution) + dims = size(sample) + + ef = convert(ExponentialFamilyDistribution, distribution) + T = ExponentialFamily.exponential_family_typetag(ef) + M = get_fisher_manifold(T, dims, getconditioner(ef)) + η = partition_point(T, dims, getnaturalparameters(ef), getconditioner(ef)) + + @test is_point(M, η, error=:error) + ef_back = convert(ExponentialFamilyDistribution, M, η) + @test getnaturalparameters(ef_back) ≈ getnaturalparameters(ef) + @test getconditioner(ef_back) == getconditioner(ef) + @test isproper(ef_back) == true + @test ExponentialFamily.exponential_family_typetag(M) == + ExponentialFamily.exponential_family_typetag(ef) + + # proper forwarding + @test ManifoldsBase.get_forwarding_type(M, inner) == + ManifoldsBase.StopForwardingType() + @test ManifoldsBase.get_forwarding_type(M, norm) == + ManifoldsBase.StopForwardingType() + + # check jacobians are inverses + natural_manifold = ManifoldsBase.decorated_manifold(M) + for _ in 1:test_points + p = rand(rng, M) + v = rand(rng, M; vector_at=p) + + v_nat = ExponentialFamilyManifolds.jacobian_manifold_to_nat(natural_manifold, v) + v_back = ExponentialFamilyManifolds.jacobian_nat_to_manifold( + natural_manifold, v_nat + ) + @test v ≈ v_back + + v_nat_2 = rand(rng, M; vector_at=p) # vectors are in same space + v_man = ExponentialFamilyManifolds.jacobian_nat_to_manifold( + natural_manifold, v_nat_2 + ) + v_nat_back = ExponentialFamilyManifolds.jacobian_manifold_to_nat( + natural_manifold, v_man + ) + @test v_nat_2 ≈ v_nat_back + end + + # Explicit fisher check for one point + p_test = rand(rng, M) + X_test = rand(rng, M; vector_at=p_test) + fisher_matrix = fisherinformation(convert(ExponentialFamilyDistribution, M, p_test)) + + # Since for Bernoulli and most simple cases without coordinate change, + # jacobian is Identity, so inner should be X' * Fisher * X. + # We check if the inner product computed by the manifold matches explicit Fisher calculation. + # But first we need to know if there is a coordinate change. + # For Bernoulli, natural params are just the vector itself, so no coordinate change. + + X_nat_test = ExponentialFamilyManifolds.jacobian_manifold_to_nat(natural_manifold, X_test) + expected_inner = dot(X_nat_test, fisher_matrix, X_nat_test) + + @test inner(M, p_test, X_test, X_test) ≈ expected_inner + + #respect fisher metric + for _ in 1:test_points + p = rand(rng, M) + v = rand(rng, M; vector_at=p) + v_norm = maximal_norm*v ./ norm(M, p, v) + v2 = rand(rng, M; vector_at=p) + @test check_geodesic(M, p, v_norm, tol=tol, error=:warn) + default_vector_transport = ManifoldsBase.default_vector_transport_method(M) + @test check_vector_transport(M, default_vector_transport, p, v_norm, v2) + end + end +end diff --git a/test/metric_manifolds/mle_metric_manifolds_setuptests.jl b/test/metric_manifolds/mle_metric_manifolds_setuptests.jl new file mode 100644 index 0000000..17a37a5 --- /dev/null +++ b/test/metric_manifolds/mle_metric_manifolds_setuptests.jl @@ -0,0 +1,74 @@ +using StableRNGs, ExponentialFamily, ManifoldsBase, LinearAlgebra +using Distributions, Random + +using Manopt +import Distributions: kldivergence, Distribution + +import ExponentialFamilyManifolds: get_fisher_manifold, partition_point +import ADTypes: AutoForwardDiff +using ManifoldDiff +import ManifoldDiff: TangentDiffBackend + +function test_mle_works( + f; + seed=42, + mle_samples=1000, + ndistributions=10, + backend_type=AutoForwardDiff(), + kl_friendly=true, + use_sectional_curvature=false, + sectional_curvature_bound=0.0, +) + rng = StableRNG(seed) + + foreach(1:ndistributions) do _ + distribution = f(rng) + ef = convert(ExponentialFamilyDistribution, distribution) + T = ExponentialFamily.exponential_family_typetag(ef) + dims = size(rand(rng, distribution)) + conditioner = getconditioner(ef) + M = get_fisher_manifold(T, dims, conditioner) + + # Generate samples from the distribution + samples = [rand(rng, distribution) for _ in 1:mle_samples] + + function cost(M, p) + ef_candidate = convert(ExponentialFamilyDistribution, M, p) + return -mean(s -> ExponentialFamily.logpdf(ef_candidate, s), samples) + end + + function grad(M, p, backend=TangentDiffBackend(backend_type)) + return ManifoldDiff.gradient(M, (p) -> cost(M, p), p, backend) + end + + if use_sectional_curvature + stepsize = DistanceOverGradients( + M; use_curvature=true, sectional_curvature_bound=sectional_curvature_max(M) + ) + else + stepsize = DistanceOverGradients() + end + # Use more iterations and tighter tolerance for better convergence + stopping_criterion = StopWhenGradientNormLess(1e-4) + p_mle = gradient_descent( + M, + cost, + grad, + rand(rng, M); + stepsize=stepsize, + stopping_criterion=stopping_criterion, + ) + ef_mle = convert(ExponentialFamilyDistribution, M, p_mle) + if kl_friendly + # Some of the distributions are using samples inside of the + # `kldivergence` function, so we need to seed the random number generator + # the `kldivergence` is defined in Distributions.jl + # also see https://github.com/JuliaStats/Distributions.jl/issues/1667 + Random.seed!(Random.default_rng(), rand(rng, UInt64)) + kl_div = kldivergence(convert(Distribution, ef_mle), distribution) + @test kl_div < 0.1 + else + @test getnaturalparameters(ef_mle) ≈ getnaturalparameters(ef) atol = 4e-1 + end + end +end diff --git a/test/metric_manifolds/normal_tests.jl b/test/metric_manifolds/normal_tests.jl new file mode 100644 index 0000000..a6886ed --- /dev/null +++ b/test/metric_manifolds/normal_tests.jl @@ -0,0 +1,14 @@ +@testitem "Check `Normal` natural manifold" begin + include("metric_manifolds_setuptests.jl") + + test_metric_manifold(tol=1e-4) do rng + return NormalMeanVariance(10randn(rng), 10rand(rng)) + end +end + +@testitem "Check `Normal` fisher manifold MLE works" begin + include("mle_metric_manifolds_setuptests.jl") + test_mle_works() do rng + return NormalMeanVariance(10randn(rng), 10rand(rng)) + end +end diff --git a/test/natural_manifolds/bernoulli_tests.jl b/test/natural_manifolds/bernoulli_tests.jl index 5c290aa..6a471df 100644 --- a/test/natural_manifolds/bernoulli_tests.jl +++ b/test/natural_manifolds/bernoulli_tests.jl @@ -8,7 +8,7 @@ end @testitem "Check MLE works for `Bernoulli`" begin include("mle_manifolds_setuptests.jl") - test_mle_works(; mle_samples=1000, ndistributions=3) do rng + test_mle_works() do rng return Bernoulli(rand(rng)) end end