Fisher manifolds

.gitignore

@@ -4,4 +4,4 @@
 /Manifest.toml
 /docs/Manifest.toml
 /docs/build/
-*Manifest.toml
+scripts/Manifest.toml

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,15 +28,16 @@ ManifoldsBase = "2.2.0"
 Random = "1.10"
 RecursiveArrayTools = "3"
 Static = "0.8, 1"
+StaticArrays = "1.9.15"
 julia = "1.10"
 
 [extras]
 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"

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

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

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