From 59377c279683d559b91063b80343b126c0f16533 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Tue, 11 Nov 2025 17:07:54 +0100 Subject: [PATCH 01/28] feat: draft WithMetric manifold --- src/ExponentialFamilyManifolds.jl | 4 ++ src/metric_manifolds/bernoulli.jl | 39 +++++++++++++++ src/natural_manifolds.jl | 2 +- src/with_metric.jl | 48 +++++++++++++++++++ test/metric_manifolds/bernoulli_tests.jl | 17 +++++++ .../metric_manifolds_setuptests.jl | 36 ++++++++++++++ test/natural_manifolds/bernoulli_tests.jl | 14 ------ 7 files changed, 145 insertions(+), 15 deletions(-) create mode 100644 src/metric_manifolds/bernoulli.jl create mode 100644 src/with_metric.jl create mode 100644 test/metric_manifolds/bernoulli_tests.jl create mode 100644 test/metric_manifolds/metric_manifolds_setuptests.jl delete mode 100644 test/natural_manifolds/bernoulli_tests.jl diff --git a/src/ExponentialFamilyManifolds.jl b/src/ExponentialFamilyManifolds.jl index 51b6fd7..bc13812 100644 --- a/src/ExponentialFamilyManifolds.jl +++ b/src/ExponentialFamilyManifolds.jl @@ -4,6 +4,7 @@ using BayesBase, ExponentialFamily, ManifoldsBase, Manifolds, Random, LinearAlge 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 +25,7 @@ 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/bernoulli.jl") end diff --git a/src/metric_manifolds/bernoulli.jl b/src/metric_manifolds/bernoulli.jl new file mode 100644 index 0000000..cd7005e --- /dev/null +++ b/src/metric_manifolds/bernoulli.jl @@ -0,0 +1,39 @@ +@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::Any, ::typeof(exp)) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::Any, ::typeof(Manifolds.exp_fused)) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::Any, ::typeof(Manifolds.geodesic)) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::Any, ::typeof(log)) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::typeof(exp)) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::typeof(Manifolds.exp_fused)) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::typeof(Manifolds.geodesic)) where {F} = _STOP +@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::typeof(log)) where {F} = _STOP + + +function _geodesic_eta(eta0, X) + t = 1 + sqrtg0 = exp.(eta0/2) / (1 .+ exp.(eta0)) + s0 = 2.0.*atan.(exp.(eta0/2)) + s = s0 .+ t .* sqrtg0 .* X + 2*log.(tan.(s./2)) +end + +# Log map (t=1): solve for X given eta1 +function _log_eta(eta0, eta1) + sqrtg0 = exp(eta0/2) / (1 + exp(eta0)) + s0 = 2*atan(exp(eta0/2)) + s1 = 2*atan(exp(eta1/2)) + (s1 - s0) / sqrtg0 +end + +ManifoldsBase.exp!(::WithMetric{F, Bernoulli, NaturalMetric}, p, X) where {F} = begin + error("I call error here!") + _geodesic_eta(p, X) +end + +ManifoldsBase.exp(::WithMetric{F, Bernoulli, NaturalMetric}, p, X) where {F} = begin + error("I call error here!") + _geodesic_eta(p, X) +end + +ManifoldsBase.log(::WithMetric{F, Bernoulli, NaturalMetric}, p, q) where {F} = begin + _log_eta(p, q) +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..7fe4a34 --- /dev/null +++ b/src/with_metric.jl @@ -0,0 +1,48 @@ +const _STOP = ManifoldsBase.StopForwardingType() +const _FWD = ManifoldsBase.SimpleForwardingType() + +abstract type MetricMode end +struct NaturalMetric <: MetricMode end # Fisher info in active coordinates + +struct WithMetric{𝔽, T, Mode<:MetricMode, M<:NaturalParametersManifold} <: AbstractDecoratorManifold{𝔽} + man::M +end +ManifoldsBase.decorated_manifold(W::WithMetric) = W.man + +# convenience constructors +with_natural_metric(M::NaturalParametersManifold{F}) where {F} = begin + 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}, ::Any, ::typeof(inner)) where {F, T} = _STOP +@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, T, NaturalMetric}, ::Any, ::typeof(norm)) where {F, T} = _STOP +@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} + ef = convert(ExponentialFamilyDistribution, M, p) + fisher_info = fisherinformation(ef) + return dot(X, fisher_info, Y) +end \ No newline at end of file diff --git a/test/metric_manifolds/bernoulli_tests.jl b/test/metric_manifolds/bernoulli_tests.jl new file mode 100644 index 0000000..a511dbf --- /dev/null +++ b/test/metric_manifolds/bernoulli_tests.jl @@ -0,0 +1,17 @@ +using Test +using Pkg + +Pkg.activate(".") + + +include("metric_manifolds_setuptests.jl") + +test_metric_manifold() do rng + return Bernoulli(rand(rng)) +end + +M = ExponentialFamilyManifolds.get_natural_manifold(Bernoulli, ()) +FM = ExponentialFamilyManifolds.with_natural_metric(M) +p = rand(FM) +right_2 = ExponentialFamilyManifolds._geodesic_eta(p, [1.0]) +@test exp(FM, p, [1]) == ExponentialFamilyManifolds._geodesic_eta(p, [1.0]) \ No newline at end of file diff --git a/test/metric_manifolds/metric_manifolds_setuptests.jl b/test/metric_manifolds/metric_manifolds_setuptests.jl new file mode 100644 index 0000000..25ba5de --- /dev/null +++ b/test/metric_manifolds/metric_manifolds_setuptests.jl @@ -0,0 +1,36 @@ +using StableRNGs, ExponentialFamily, ManifoldsBase, LinearAlgebra +using Distributions + +import ExponentialFamilyManifolds: get_natural_manifold, with_natural_metric, partition_point + +function test_metric_manifold(f; seed=42, ndistributions=100, test_points = 5) + 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) + NM = get_natural_manifold(T, dims, getconditioner(ef)) + M = with_natural_metric(NM) + η = 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() + + #respect fisher metric + for _ in 1:test_points + @test check_geodesic(M, rand(rng, M), rand(rng, M), error = :error) + end + end +end diff --git a/test/natural_manifolds/bernoulli_tests.jl b/test/natural_manifolds/bernoulli_tests.jl deleted file mode 100644 index 5c290aa..0000000 --- a/test/natural_manifolds/bernoulli_tests.jl +++ /dev/null @@ -1,14 +0,0 @@ -@testitem "Check `Bernoulli` natural manifold" begin - include("natural_manifolds_setuptests.jl") - - test_natural_manifold() do rng - return Bernoulli(rand(rng)) - end -end - -@testitem "Check MLE works for `Bernoulli`" begin - include("mle_manifolds_setuptests.jl") - test_mle_works(; mle_samples=1000, ndistributions=3) do rng - return Bernoulli(rand(rng)) - end -end From 5a1cc613099d9c29d545538f32914edaa0c4868e Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 14 Nov 2025 11:37:01 +0100 Subject: [PATCH 02/28] test: add missing test for Bernoulli natural manifold --- test/natural_manifolds/bernoulli_tests.jl | 14 ++++++++++++++ 1 file changed, 14 insertions(+) create mode 100644 test/natural_manifolds/bernoulli_tests.jl diff --git a/test/natural_manifolds/bernoulli_tests.jl b/test/natural_manifolds/bernoulli_tests.jl new file mode 100644 index 0000000..6a471df --- /dev/null +++ b/test/natural_manifolds/bernoulli_tests.jl @@ -0,0 +1,14 @@ +@testitem "Check `Bernoulli` natural manifold" begin + include("natural_manifolds_setuptests.jl") + + test_natural_manifold() do rng + return Bernoulli(rand(rng)) + end +end + +@testitem "Check MLE works for `Bernoulli`" begin + include("mle_manifolds_setuptests.jl") + test_mle_works() do rng + return Bernoulli(rand(rng)) + end +end From f882c81b5e26e20b9485e9f01ab50ce30e7e60de Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 14 Nov 2025 11:37:46 +0100 Subject: [PATCH 03/28] fix: add metric manifold tests --- test/metric_manifolds/bernoulli_tests.jl | 23 +++++-------------- .../metric_manifolds_setuptests.jl | 7 ++++-- 2 files changed, 11 insertions(+), 19 deletions(-) diff --git a/test/metric_manifolds/bernoulli_tests.jl b/test/metric_manifolds/bernoulli_tests.jl index a511dbf..c247833 100644 --- a/test/metric_manifolds/bernoulli_tests.jl +++ b/test/metric_manifolds/bernoulli_tests.jl @@ -1,17 +1,6 @@ -using Test -using Pkg - -Pkg.activate(".") - - -include("metric_manifolds_setuptests.jl") - -test_metric_manifold() do rng - return Bernoulli(rand(rng)) -end - -M = ExponentialFamilyManifolds.get_natural_manifold(Bernoulli, ()) -FM = ExponentialFamilyManifolds.with_natural_metric(M) -p = rand(FM) -right_2 = ExponentialFamilyManifolds._geodesic_eta(p, [1.0]) -@test exp(FM, p, [1]) == ExponentialFamilyManifolds._geodesic_eta(p, [1.0]) \ No newline at end of file +@testitem "Check `Bernoulli` fisher manifold" begin + include("metric_manifolds_setuptests.jl") + test_metric_manifold() do rng + return Bernoulli(rand(rng)) + end +end \ No newline at end of file diff --git a/test/metric_manifolds/metric_manifolds_setuptests.jl b/test/metric_manifolds/metric_manifolds_setuptests.jl index 25ba5de..9bf91c5 100644 --- a/test/metric_manifolds/metric_manifolds_setuptests.jl +++ b/test/metric_manifolds/metric_manifolds_setuptests.jl @@ -3,7 +3,7 @@ using Distributions import ExponentialFamilyManifolds: get_natural_manifold, with_natural_metric, partition_point -function test_metric_manifold(f; seed=42, ndistributions=100, test_points = 5) +function test_metric_manifold(f; seed=42, ndistributions=100, test_points = 10) rng = StableRNG(seed) foreach(1:ndistributions) do _ @@ -30,7 +30,10 @@ function test_metric_manifold(f; seed=42, ndistributions=100, test_points = 5) #respect fisher metric for _ in 1:test_points - @test check_geodesic(M, rand(rng, M), rand(rng, M), error = :error) + p = rand(rng, M) + v = rand(rng, M; vector_at = p) + v_norm = 0.42*v./norm(M, p, v) + @test check_geodesic(M, p, v_norm, tol=1e-3, error=:warn) end end end From a4a0e1fe4c797692162907a1c6524287b52be792 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 14 Nov 2025 11:39:46 +0100 Subject: [PATCH 04/28] style: :art: --- src/with_metric.jl | 45 +++++++++++-------- test/metric_manifolds/bernoulli_tests.jl | 2 +- .../metric_manifolds_setuptests.jl | 18 +++++--- 3 files changed, 39 insertions(+), 26 deletions(-) diff --git a/src/with_metric.jl b/src/with_metric.jl index 7fe4a34..6a9cc31 100644 --- a/src/with_metric.jl +++ b/src/with_metric.jl @@ -1,48 +1,57 @@ const _STOP = ManifoldsBase.StopForwardingType() -const _FWD = ManifoldsBase.SimpleForwardingType() +const _FWD = ManifoldsBase.SimpleForwardingType() abstract type MetricMode end -struct NaturalMetric <: MetricMode end # Fisher info in active coordinates +struct NaturalMetric <: MetricMode end # Fisher info in active coordinates -struct WithMetric{𝔽, T, Mode<:MetricMode, M<:NaturalParametersManifold} <: AbstractDecoratorManifold{𝔽} +struct WithMetric{𝔽,T,Mode<:MetricMode,M<:NaturalParametersManifold} <: + AbstractDecoratorManifold{𝔽} man::M end ManifoldsBase.decorated_manifold(W::WithMetric) = W.man # convenience constructors -with_natural_metric(M::NaturalParametersManifold{F}) where {F} = begin +function with_natural_metric(M::NaturalParametersManifold{F}) where {F} ef_typetag = ExponentialFamily.exponential_family_typetag(M) - WithMetric{F, ef_typetag, NaturalMetric, typeof(M)}(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}, ::Any, ::typeof(inner)) where {F, T} = _STOP -@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, T, NaturalMetric}, ::Any, ::typeof(norm)) where {F, T} = _STOP -@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 +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,T,NaturalMetric}, ::Any, ::typeof(inner) +) where {F,T} = _STOP +@inline ManifoldsBase.get_forwarding_type( + ::WithMetric{F,T,NaturalMetric}, ::Any, ::typeof(norm) +) where {F,T} = _STOP +@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 +ExponentialFamily.exponential_family_typetag(::WithMetric{F,T}) where {F,T} = T -function Base.convert( - ::Type{ExponentialFamilyDistribution}, M::WithMetric, p -) +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 + 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} +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} +function ManifoldsBase.inner(M::WithMetric{F,T,NaturalMetric}, p, X, Y) where {F,T} ef = convert(ExponentialFamilyDistribution, M, p) fisher_info = fisherinformation(ef) return dot(X, fisher_info, Y) -end \ No newline at end of file +end diff --git a/test/metric_manifolds/bernoulli_tests.jl b/test/metric_manifolds/bernoulli_tests.jl index c247833..b7a7e3a 100644 --- a/test/metric_manifolds/bernoulli_tests.jl +++ b/test/metric_manifolds/bernoulli_tests.jl @@ -3,4 +3,4 @@ test_metric_manifold() do rng return Bernoulli(rand(rng)) end -end \ No newline at end of file +end diff --git a/test/metric_manifolds/metric_manifolds_setuptests.jl b/test/metric_manifolds/metric_manifolds_setuptests.jl index 9bf91c5..ce9587b 100644 --- a/test/metric_manifolds/metric_manifolds_setuptests.jl +++ b/test/metric_manifolds/metric_manifolds_setuptests.jl @@ -1,9 +1,10 @@ using StableRNGs, ExponentialFamily, ManifoldsBase, LinearAlgebra using Distributions -import ExponentialFamilyManifolds: get_natural_manifold, with_natural_metric, partition_point +import ExponentialFamilyManifolds: + get_natural_manifold, with_natural_metric, partition_point -function test_metric_manifold(f; seed=42, ndistributions=100, test_points = 10) +function test_metric_manifold(f; seed=42, ndistributions=100, test_points=10) rng = StableRNG(seed) foreach(1:ndistributions) do _ @@ -22,17 +23,20 @@ function test_metric_manifold(f; seed=42, ndistributions=100, test_points = 10) @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) + @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() + @test ManifoldsBase.get_forwarding_type(M, inner) == + ManifoldsBase.StopForwardingType() + @test ManifoldsBase.get_forwarding_type(M, norm) == + ManifoldsBase.StopForwardingType() #respect fisher metric for _ in 1:test_points p = rand(rng, M) - v = rand(rng, M; vector_at = p) - v_norm = 0.42*v./norm(M, p, v) + v = rand(rng, M; vector_at=p) + v_norm = 0.42*v ./ norm(M, p, v) @test check_geodesic(M, p, v_norm, tol=1e-3, error=:warn) end end From d18a8188f8d38b43e8dca1217e35faec8d53a78c Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 14 Nov 2025 11:44:31 +0100 Subject: [PATCH 05/28] test(fix): setup generic mle works for Fisher manifolds --- src/metric_manifolds/bernoulli.jl | 140 ++++++++++++++++++----- test/metric_manifolds/bernoulli_tests.jl | 7 ++ 2 files changed, 119 insertions(+), 28 deletions(-) diff --git a/src/metric_manifolds/bernoulli.jl b/src/metric_manifolds/bernoulli.jl index cd7005e..bbf99c3 100644 --- a/src/metric_manifolds/bernoulli.jl +++ b/src/metric_manifolds/bernoulli.jl @@ -1,39 +1,123 @@ -@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::Any, ::typeof(exp)) where {F} = _STOP -@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::Any, ::typeof(Manifolds.exp_fused)) where {F} = _STOP -@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::Any, ::typeof(Manifolds.geodesic)) where {F} = _STOP -@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::Any, ::typeof(log)) where {F} = _STOP -@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::typeof(exp)) where {F} = _STOP -@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::typeof(Manifolds.exp_fused)) where {F} = _STOP -@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::typeof(Manifolds.geodesic)) where {F} = _STOP -@inline ManifoldsBase.get_forwarding_type(::WithMetric{F, Bernoulli, NaturalMetric}, ::typeof(log)) where {F} = _STOP +function geodesic_exact(t, x0, v0) + # Compute metric at initial point + g0 = exp(x0) / (1 + exp(x0))^2 + # Energy constant + E = g0 * v0^2 -function _geodesic_eta(eta0, X) - t = 1 - sqrtg0 = exp.(eta0/2) / (1 .+ exp.(eta0)) - s0 = 2.0.*atan.(exp.(eta0/2)) - s = s0 .+ t .* sqrtg0 .* X - 2*log.(tan.(s./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(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 + +# Exponential map implementations - these override the default forwarding behavior +function ManifoldsBase.exp!(::WithMetric{F,Bernoulli,NaturalMetric}, q, p, X) where {F} + x, _ = geodesic_exact(1, p[1], X[1]) + q .= x + return q +end + +function ManifoldsBase.exp(::WithMetric{F,Bernoulli,NaturalMetric}, p, X) where {F} + x, _ = geodesic_exact(1, p[1], X[1]) + return [x] +end + +function ManifoldsBase.exp_fused!( + ::WithMetric{F,Bernoulli,NaturalMetric}, q, p, X, t::Number +) where {F} + x, _ = geodesic_exact(1, p[1], X[1]) + q .= x + return q +end + +function ManifoldsBase.exp_fused( + ::WithMetric{F,Bernoulli,NaturalMetric}, p, X, t::Number +) where {F} + x, _ = geodesic_exact(t, p[1], X[1]) + return [x] +end + +function ManifoldsBase.log!(::WithMetric{F,Bernoulli,NaturalMetric}, X, p, q) where {F} + X .= log_map(p[1], q[1]) + return X +end + +function ManifoldsBase.log(::WithMetric{F,Bernoulli,NaturalMetric}, p, q) where {F} + return log_map(p[1], q[1]) +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) + X .* (s_p ./ s_q) +end + +function ManifoldsBase.parallel_transport_to!( + ::WithMetric{F,Bernoulli,NaturalMetric}, Y, p, X, q; kwargs... +) where {F} + # error("Call here") + s_p = _sqrtg_eta(p) + s_q = _sqrtg_eta(q) + @. Y = X * (s_p / s_q) + Y end -# Log map (t=1): solve for X given eta1 -function _log_eta(eta0, eta1) - sqrtg0 = exp(eta0/2) / (1 + exp(eta0)) - s0 = 2*atan(exp(eta0/2)) - s1 = 2*atan(exp(eta1/2)) - (s1 - s0) / sqrtg0 +# Parallel transport in a direction +function ManifoldsBase.parallel_transport_direction( + M::WithMetric{F,Bernoulli,NaturalMetric}, p, X, d; kwargs... +) where {F} + q = exp(M, p, d) + return parallel_transport_to(M, p, X, q) end -ManifoldsBase.exp!(::WithMetric{F, Bernoulli, NaturalMetric}, p, X) where {F} = begin - error("I call error here!") - _geodesic_eta(p, X) +function ManifoldsBase.parallel_transport_direction!( + M::WithMetric{F,Bernoulli,NaturalMetric}, Y, p, X, d; kwargs... +) where {F} + q = exp(M, p, d) + return parallel_transport_to!(M, Y, p, X, q) end -ManifoldsBase.exp(::WithMetric{F, Bernoulli, NaturalMetric}, p, X) where {F} = begin - error("I call error here!") - _geodesic_eta(p, X) +# Vector transport with ParallelTransport method - this is what check_geodesic uses +function ManifoldsBase.vector_transport_to( + M::WithMetric{F,Bernoulli,NaturalMetric}, p, X, q, ::ManifoldsBase.ParallelTransport +) where {F} + return parallel_transport_to(M, p, X, q) end -ManifoldsBase.log(::WithMetric{F, Bernoulli, NaturalMetric}, p, q) where {F} = begin - _log_eta(p, q) +function ManifoldsBase.vector_transport_to!( + M::WithMetric{F,Bernoulli,NaturalMetric}, Y, p, X, q, ::ManifoldsBase.ParallelTransport +) where {F} + return parallel_transport_to!(M, Y, p, X, q) end diff --git a/test/metric_manifolds/bernoulli_tests.jl b/test/metric_manifolds/bernoulli_tests.jl index b7a7e3a..75f5e82 100644 --- a/test/metric_manifolds/bernoulli_tests.jl +++ b/test/metric_manifolds/bernoulli_tests.jl @@ -4,3 +4,10 @@ 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 From 46c1e709121298c33bf225aa179dfa305b41c56c Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 14 Nov 2025 11:45:20 +0100 Subject: [PATCH 06/28] test(fix): add generic test mle for metric manifolds --- .../mle_metric_manifolds_setuptests.jl | 58 +++++++++++++++++++ 1 file changed, 58 insertions(+) create mode 100644 test/metric_manifolds/mle_metric_manifolds_setuptests.jl 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..1b0ff41 --- /dev/null +++ b/test/metric_manifolds/mle_metric_manifolds_setuptests.jl @@ -0,0 +1,58 @@ +using StableRNGs, ExponentialFamily, ManifoldsBase, LinearAlgebra +using Distributions, Random + +using Manopt +import Distributions: kldivergence, Distribution + +import ExponentialFamilyManifolds: get_natural_manifold, partition_point, with_natural_metric +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, +) + 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) + NM = get_natural_manifold(T, dims, conditioner) + M = with_natural_metric(NM) + + # 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 + + stepsize = DistanceOverGradients() + p_mle = gradient_descent(M, cost, grad, rand(rng, M); stepsize=stepsize) + 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 From 700da1099933d8778c02c50cb0c14c4cc2c2ffcd Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 14 Nov 2025 11:46:37 +0100 Subject: [PATCH 07/28] ci(fix): no need to store scripts manifiest --- scripts/Manifest.toml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/scripts/Manifest.toml b/scripts/Manifest.toml index a868251..61eafaf 100644 --- a/scripts/Manifest.toml +++ b/scripts/Manifest.toml @@ -1,6 +1,6 @@ # This file is machine-generated - editing it directly is not advised -julia_version = "1.10.9" +julia_version = "1.11.7" manifest_format = "2.0" project_hash = "93e024a7bf6c9615a79c4b80b3a014b0f1f744c9" @@ -19,6 +19,7 @@ version = "0.9.1" [[deps.Dates]] deps = ["Printf"] uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" +version = "1.11.0" [[deps.Glob]] git-tree-sha1 = "97285bbd5230dd766e9ef6749b80fc617126d496" @@ -38,6 +39,7 @@ version = "0.4.10" [[deps.Logging]] uuid = "56ddb016-857b-54e1-b83d-db4d58db5568" +version = "1.11.0" [[deps.PrecompileTools]] deps = ["Preferences"] @@ -54,6 +56,7 @@ version = "1.5.0" [[deps.Printf]] deps = ["Unicode"] uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7" +version = "1.11.0" [[deps.TOML]] deps = ["Dates"] @@ -67,3 +70,4 @@ version = "1.0.2" [[deps.Unicode]] uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5" +version = "1.11.0" From ad0e24cb99a2f5e88b2fe905e62752196f457952 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 14 Nov 2025 11:49:03 +0100 Subject: [PATCH 08/28] ci(fix): git ignore scripts/Manifest.toml --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 95731a5..c31bdf3 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ /Manifest.toml /docs/Manifest.toml /docs/build/ +scripts/Manifest.toml \ No newline at end of file From 3e435510a7ca7fc6779cf26e3f191477f9f4a487 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Fri, 14 Nov 2025 11:50:30 +0100 Subject: [PATCH 09/28] fix: alphabetic reodering in toml --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 7fa596d..3209e84 100644 --- a/Project.toml +++ b/Project.toml @@ -34,8 +34,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" From b5a99a667993f8fee0d029c334cc75db77f57c62 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Tue, 18 Nov 2025 14:49:25 +0100 Subject: [PATCH 10/28] fix: friendly constructor get_fisher_manifold --- src/with_metric.jl | 7 ++++++- test/metric_manifolds/bernoulli_tests.jl | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/with_metric.jl b/src/with_metric.jl index 6a9cc31..e033c84 100644 --- a/src/with_metric.jl +++ b/src/with_metric.jl @@ -10,7 +10,12 @@ struct WithMetric{𝔽,T,Mode<:MetricMode,M<:NaturalParametersManifold} <: end ManifoldsBase.decorated_manifold(W::WithMetric) = W.man -# convenience constructors +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) diff --git a/test/metric_manifolds/bernoulli_tests.jl b/test/metric_manifolds/bernoulli_tests.jl index 75f5e82..3ac51d4 100644 --- a/test/metric_manifolds/bernoulli_tests.jl +++ b/test/metric_manifolds/bernoulli_tests.jl @@ -1,6 +1,6 @@ @testitem "Check `Bernoulli` fisher manifold" begin include("metric_manifolds_setuptests.jl") - test_metric_manifold() do rng + test_metric_manifold(maximal_norm=0.4) do rng return Bernoulli(rand(rng)) end end From fb95764f913e2c655c76f273f50e8928fc45e306 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Tue, 18 Nov 2025 14:49:58 +0100 Subject: [PATCH 11/28] fix(test): use get_fisher_manifold constructor --- test/metric_manifolds/metric_manifolds_setuptests.jl | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/test/metric_manifolds/metric_manifolds_setuptests.jl b/test/metric_manifolds/metric_manifolds_setuptests.jl index ce9587b..0efcb06 100644 --- a/test/metric_manifolds/metric_manifolds_setuptests.jl +++ b/test/metric_manifolds/metric_manifolds_setuptests.jl @@ -2,9 +2,9 @@ using StableRNGs, ExponentialFamily, ManifoldsBase, LinearAlgebra using Distributions import ExponentialFamilyManifolds: - get_natural_manifold, with_natural_metric, partition_point + get_fisher_manifold, partition_point -function test_metric_manifold(f; seed=42, ndistributions=100, test_points=10) +function test_metric_manifold(f; seed=42, ndistributions=100, test_points=10, maximal_norm=1) rng = StableRNG(seed) foreach(1:ndistributions) do _ @@ -14,8 +14,7 @@ function test_metric_manifold(f; seed=42, ndistributions=100, test_points=10) ef = convert(ExponentialFamilyDistribution, distribution) T = ExponentialFamily.exponential_family_typetag(ef) - NM = get_natural_manifold(T, dims, getconditioner(ef)) - M = with_natural_metric(NM) + M = get_fisher_manifold(T, dims, getconditioner(ef)) η = partition_point(T, dims, getnaturalparameters(ef), getconditioner(ef)) @test is_point(M, η, error=:error) @@ -36,7 +35,7 @@ function test_metric_manifold(f; seed=42, ndistributions=100, test_points=10) for _ in 1:test_points p = rand(rng, M) v = rand(rng, M; vector_at=p) - v_norm = 0.42*v ./ norm(M, p, v) + v_norm = maximal_norm*v ./ norm(M, p, v) @test check_geodesic(M, p, v_norm, tol=1e-3, error=:warn) end end From f3f1686f634d17c1c0b09473b4d9ad96fda81faf Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Tue, 18 Nov 2025 16:39:30 +0100 Subject: [PATCH 12/28] test: add univariate normal metric manifold tests --- Project.toml | 2 ++ src/metric_manifolds/bernoulli.jl | 16 ++++++++-------- .../metric_manifolds_setuptests.jl | 4 ++-- .../mle_metric_manifolds_setuptests.jl | 5 ++--- test/metric_manifolds/normal_tests.jl | 14 ++++++++++++++ 5 files changed, 28 insertions(+), 13 deletions(-) create mode 100644 test/metric_manifolds/normal_tests.jl diff --git a/Project.toml b/Project.toml index 3209e84..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] diff --git a/src/metric_manifolds/bernoulli.jl b/src/metric_manifolds/bernoulli.jl index bbf99c3..6c90821 100644 --- a/src/metric_manifolds/bernoulli.jl +++ b/src/metric_manifolds/bernoulli.jl @@ -1,4 +1,4 @@ -function geodesic_exact(t, x0, v0) +function _geodesic_exact_bernoulli(t, x0, v0) # Compute metric at initial point g0 = exp(x0) / (1 + exp(x0))^2 @@ -21,7 +21,7 @@ function geodesic_exact(t, x0, v0) return x, v end -function log_map(x0, x1) +function _log_map_bernoulli(x0, x1) C = 2 * atan(exp(x0/2)) C1 = 2 * atan(exp(x1/2)) @@ -40,20 +40,20 @@ end # Exponential map implementations - these override the default forwarding behavior function ManifoldsBase.exp!(::WithMetric{F,Bernoulli,NaturalMetric}, q, p, X) where {F} - x, _ = geodesic_exact(1, p[1], X[1]) + x, _ = _geodesic_exact_bernoulli(1, p[1], X[1]) q .= x return q end function ManifoldsBase.exp(::WithMetric{F,Bernoulli,NaturalMetric}, p, X) where {F} - x, _ = geodesic_exact(1, p[1], X[1]) + x, _ = _geodesic_exact_bernoulli(1, p[1], X[1]) return [x] end function ManifoldsBase.exp_fused!( ::WithMetric{F,Bernoulli,NaturalMetric}, q, p, X, t::Number ) where {F} - x, _ = geodesic_exact(1, p[1], X[1]) + x, _ = _geodesic_exact_bernoulli(1, p[1], X[1]) q .= x return q end @@ -61,17 +61,17 @@ end function ManifoldsBase.exp_fused( ::WithMetric{F,Bernoulli,NaturalMetric}, p, X, t::Number ) where {F} - x, _ = geodesic_exact(t, p[1], X[1]) + x, _ = _geodesic_exact_bernoulli(t, p[1], X[1]) return [x] end function ManifoldsBase.log!(::WithMetric{F,Bernoulli,NaturalMetric}, X, p, q) where {F} - X .= log_map(p[1], q[1]) + X .= _log_map_bernoulli(p[1], q[1]) return X end function ManifoldsBase.log(::WithMetric{F,Bernoulli,NaturalMetric}, p, q) where {F} - return log_map(p[1], q[1]) + return _log_map_bernoulli(p[1], q[1]) end _sqrtg_eta(η) = 0.5 .* sech.(η ./ 2) diff --git a/test/metric_manifolds/metric_manifolds_setuptests.jl b/test/metric_manifolds/metric_manifolds_setuptests.jl index 0efcb06..cdc3d2a 100644 --- a/test/metric_manifolds/metric_manifolds_setuptests.jl +++ b/test/metric_manifolds/metric_manifolds_setuptests.jl @@ -4,7 +4,7 @@ using Distributions import ExponentialFamilyManifolds: get_fisher_manifold, partition_point -function test_metric_manifold(f; seed=42, ndistributions=100, test_points=10, maximal_norm=1) +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 _ @@ -36,7 +36,7 @@ function test_metric_manifold(f; seed=42, ndistributions=100, test_points=10, ma p = rand(rng, M) v = rand(rng, M; vector_at=p) v_norm = maximal_norm*v ./ norm(M, p, v) - @test check_geodesic(M, p, v_norm, tol=1e-3, error=:warn) + @test check_geodesic(M, p, v_norm, tol=tol, error=:warn) end end end diff --git a/test/metric_manifolds/mle_metric_manifolds_setuptests.jl b/test/metric_manifolds/mle_metric_manifolds_setuptests.jl index 1b0ff41..e6f52d1 100644 --- a/test/metric_manifolds/mle_metric_manifolds_setuptests.jl +++ b/test/metric_manifolds/mle_metric_manifolds_setuptests.jl @@ -4,7 +4,7 @@ using Distributions, Random using Manopt import Distributions: kldivergence, Distribution -import ExponentialFamilyManifolds: get_natural_manifold, partition_point, with_natural_metric +import ExponentialFamilyManifolds: get_fisher_manifold, partition_point import ADTypes: AutoForwardDiff using ManifoldDiff import ManifoldDiff: TangentDiffBackend @@ -25,8 +25,7 @@ function test_mle_works( T = ExponentialFamily.exponential_family_typetag(ef) dims = size(rand(rng, distribution)) conditioner = getconditioner(ef) - NM = get_natural_manifold(T, dims, conditioner) - M = with_natural_metric(NM) + M = get_fisher_manifold(T, dims, conditioner) # Generate samples from the distribution samples = [rand(rng, distribution) for _ in 1:mle_samples] diff --git a/test/metric_manifolds/normal_tests.jl b/test/metric_manifolds/normal_tests.jl new file mode 100644 index 0000000..329eb63 --- /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-2, maximal_norm=0.11 ) do rng + return NormalMeanVariance(10randn(rng), rand(rng)+1) + 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 \ No newline at end of file From d1acf17037f552fdf65ce1937a4dccd94a1c8e5a Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Tue, 18 Nov 2025 16:45:50 +0100 Subject: [PATCH 13/28] feaT: initial draft univariate normal manifold --- src/ExponentialFamilyManifolds.jl | 3 +- src/metric_manifolds/normal.jl | 424 ++++++++++++++++++++++++++ test/metric_manifolds/normal_tests.jl | 2 +- 3 files changed, 427 insertions(+), 2 deletions(-) create mode 100644 src/metric_manifolds/normal.jl diff --git a/src/ExponentialFamilyManifolds.jl b/src/ExponentialFamilyManifolds.jl index bc13812..ffd8731 100644 --- a/src/ExponentialFamilyManifolds.jl +++ b/src/ExponentialFamilyManifolds.jl @@ -1,6 +1,6 @@ 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") @@ -28,4 +28,5 @@ include("natural_manifolds/wishart.jl") # metric Manifolds include("metric_manifolds/bernoulli.jl") +include("metric_manifolds/normal.jl") end diff --git a/src/metric_manifolds/normal.jl b/src/metric_manifolds/normal.jl new file mode 100644 index 0000000..cdd2bae --- /dev/null +++ b/src/metric_manifolds/normal.jl @@ -0,0 +1,424 @@ +using StaticArrays +using RecursiveArrayTools: ArrayPartition + +# ============================================================================ +# 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]] + +# --- Basic conversions: natural ↔ mean/variance ↔ half-plane --- + +# (η₁, η₂) -> (μ, σ²) +@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 + Jη_xy = _J_eta_to_xy(η0) + V0 = Jη_xy * Xη0 + Jxy_u = _J_hp_to_hyp(xy0) + W0 = Jxy_u * V0 + + # 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) + Ju_xy = _J_hyp_to_hp(up) + Vp = Ju_xy * ξ + Jxy_η = _J_xy_to_eta(xyp) + Xη = Jxy_η * Vp + + 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 + +# ============================================================================ +# ManifoldsBase method overrides +# ============================================================================ + +# Exponential map implementations + +function ManifoldsBase.exp!( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + q, + p, + X, +) where {F} + p_t = geodesic_exact(1, p, X) + q[1] = p_t[1] + q[2] = p_t[2] + return q +end + +function ManifoldsBase.exp( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + p, + X, +) where {F} + p_t = geodesic_exact(1, p, X) + return ArrayPartition([p_t[1]], [p_t[2]]) +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.exp_fused( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + p, + X, + t::Number, +) where {F} + p_t = geodesic_exact(t, p, X) + return ArrayPartition([p_t[1]], [p_t[2]]) +end + +# Log map implementations + +function ManifoldsBase.log!( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + X, + p, + q, +) where {F} + X .= log_map(p, q) + return X +end + +function ManifoldsBase.log( + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, + p, + q, +) where {F} + X = log_map(p, q) + return ArrayPartition([X[1]], [X[2]]) +end + +# Parallel transport in natural parameters + +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 + Jη_xy = _J_eta_to_xy(ηp) + Vp = Jη_xy * Xη + Jxy_u = _J_hp_to_hyp(xyp) + Wp = Jxy_u * Vp + + # parallel transport on hyperboloid + Wq = _hyperbolic_parallel_transport(up, uq, Wp) + + # pull back to η at q + Ju_xy = _J_hyp_to_hp(uq) + Vq = Ju_xy * Wq + Jxy_η = _J_xy_to_eta(xyq) + Xηq = Jxy_η * Vq + + 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 + +# Direction-based PT using exp + +function ManifoldsBase.parallel_transport_direction( + M::WithMetric{F,NormalMeanVariance,NaturalMetric}, + p, + X, + d; + kwargs..., +) where {F} + q = ManifoldsBase.exp(M, p, d) + return ManifoldsBase.parallel_transport_to(M, p, X, q; kwargs...) +end + +function ManifoldsBase.parallel_transport_direction!( + M::WithMetric{F,NormalMeanVariance,NaturalMetric}, + Y, + p, + X, + d; + kwargs..., +) where {F} + q = ManifoldsBase.exp(M, p, d) + return ManifoldsBase.parallel_transport_to!(M, Y, p, X, q; kwargs...) +end + +# vector_transport_to aliases ParallelTransport, like in your Bernoulli code + +function ManifoldsBase.vector_transport_to( + M::WithMetric{F,NormalMeanVariance,NaturalMetric}, + p, + X, + q, + ::ManifoldsBase.ParallelTransport, +) where {F} + return ManifoldsBase.parallel_transport_to(M, p, X, q) +end + +function ManifoldsBase.vector_transport_to!( + M::WithMetric{F,NormalMeanVariance,NaturalMetric}, + Y, + p, + X, + q, + ::ManifoldsBase.ParallelTransport, +) where {F} + return ManifoldsBase.parallel_transport_to!(M, Y, p, X, q) +end \ No newline at end of file diff --git a/test/metric_manifolds/normal_tests.jl b/test/metric_manifolds/normal_tests.jl index 329eb63..48cc2f9 100644 --- a/test/metric_manifolds/normal_tests.jl +++ b/test/metric_manifolds/normal_tests.jl @@ -1,7 +1,7 @@ @testitem "Check `Normal` natural manifold" begin include("metric_manifolds_setuptests.jl") - test_metric_manifold(tol=1e-2, maximal_norm=0.11 ) do rng + test_metric_manifold(tol=1e-5, maximal_norm=0.3 ) do rng return NormalMeanVariance(10randn(rng), rand(rng)+1) end end From 44217eda6fa88901e17fec61d51fd003e6882fe5 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Tue, 18 Nov 2025 22:22:24 +0100 Subject: [PATCH 14/28] test: temporary quick runner for the tests --- run_normal_tests.jl | 53 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100755 run_normal_tests.jl diff --git a/run_normal_tests.jl b/run_normal_tests.jl new file mode 100755 index 0000000..317557d --- /dev/null +++ b/run_normal_tests.jl @@ -0,0 +1,53 @@ +#!/usr/bin/env julia + +# Minimal script to run Normal manifold tests outside the test suite + +using Pkg +Pkg.activate(".") + +# Load the package +using ExponentialFamilyManifolds + +# Load test dependencies +using Test +using StableRNGs +using ExponentialFamily +using ManifoldsBase +using LinearAlgebra +using Distributions +using Random +using Manopt +using ManifoldDiff +import ADTypes: AutoForwardDiff +import Distributions: kldivergence, Distribution +import ExponentialFamilyManifolds: get_fisher_manifold, partition_point +import ManifoldDiff: TangentDiffBackend + +# Include setup test functions +include("test/metric_manifolds/metric_manifolds_setuptests.jl") +include("test/metric_manifolds/mle_metric_manifolds_setuptests.jl") + +println("=" ^ 70) +println("Testing Normal metric manifold") +println("=" ^ 70) + +@testset "Normal metric manifold" begin + test_metric_manifold(tol=1e-5, maximal_norm=0.3) do rng + return NormalMeanVariance(10randn(rng), rand(rng) + 1) + end +end + +println("\n" * "=" ^ 70) +println("Testing Normal fisher manifold MLE") +println("=" ^ 70) + +@testset "Normal fisher manifold MLE" begin + test_mle_works() do rng + return NormalMeanVariance(10randn(rng), 10rand(rng)) + end +end + +println("\n" * "=" ^ 70) +println("All tests completed!") +println("=" ^ 70) + From 3e0d452edc0ff4fcd0ed8a9af6070b1f8bd15e0e Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Tue, 18 Nov 2025 22:46:00 +0100 Subject: [PATCH 15/28] fix: numerical improvment --- src/metric_manifolds/normal.jl | 14 ++++++++------ test/metric_manifolds/normal_tests.jl | 2 +- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/src/metric_manifolds/normal.jl b/src/metric_manifolds/normal.jl index cdd2bae..fef9f39 100644 --- a/src/metric_manifolds/normal.jl +++ b/src/metric_manifolds/normal.jl @@ -197,10 +197,11 @@ function geodesic_exact_eta( 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 + W0 = (Jxy_u * V0) / sqrt(T(2)) # geodesic on hyperboloid u_t = _hyperbolic_exp(u0, W0, T(t)) @@ -237,10 +238,11 @@ function log_map_eta(ηp::SVector{2,T}, ηq::SVector{2,T}) where {T<:Real} ξ = _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 + Xη = (Jxy_η * Vp) * sqrt(T(2)) return Xη end @@ -344,20 +346,20 @@ function ManifoldsBase.parallel_transport_to( up = _hp_to_hyp(xyp) uq = _hp_to_hyp(xyq) - # push X to hyperboloid + # 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 + Wp = (Jxy_u * Vp) / sqrt(2.0) # parallel transport on hyperboloid Wq = _hyperbolic_parallel_transport(up, uq, Wp) - # pull back to η at q + # 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 + Xηq = (Jxy_η * Vq) * sqrt(2.0) Xq = _eta_to_X(Xηq) return ArrayPartition([Xq[1]], [Xq[2]]) diff --git a/test/metric_manifolds/normal_tests.jl b/test/metric_manifolds/normal_tests.jl index 48cc2f9..815913c 100644 --- a/test/metric_manifolds/normal_tests.jl +++ b/test/metric_manifolds/normal_tests.jl @@ -1,7 +1,7 @@ @testitem "Check `Normal` natural manifold" begin include("metric_manifolds_setuptests.jl") - test_metric_manifold(tol=1e-5, maximal_norm=0.3 ) do rng + test_metric_manifold(tol=1e-4, maximal_norm=0.3 ) do rng return NormalMeanVariance(10randn(rng), rand(rng)+1) end end From 8e70ba4bd5378aa56691e57c6175a0fdfdb177c6 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Tue, 18 Nov 2025 23:03:59 +0100 Subject: [PATCH 16/28] fix: move normal specific methods into normal folder --- src/metric_manifolds/normal.jl | 17 +++++++++++++++++ src/with_metric.jl | 1 + .../mle_metric_manifolds_setuptests.jl | 16 ++++++++++++++-- test/metric_manifolds/normal_tests.jl | 4 ++-- 4 files changed, 34 insertions(+), 4 deletions(-) diff --git a/src/metric_manifolds/normal.jl b/src/metric_manifolds/normal.jl index fef9f39..44a0be2 100644 --- a/src/metric_manifolds/normal.jl +++ b/src/metric_manifolds/normal.jl @@ -423,4 +423,21 @@ function ManifoldsBase.vector_transport_to!( ::ManifoldsBase.ParallelTransport, ) where {F} return ManifoldsBase.parallel_transport_to!(M, Y, p, X, q) +end + +function ManifoldsBase.inner(M::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, Y) where {F} + # For NormalMeanVariance, manifold coords are (η₁, λ) but Fisher matrix is in (η₁, η₂) + # Transform tangent vectors: (dη₁, dλ) -> (dη₁, dη₂) where dη₂ = -dλ + Xη = [X[1], -X[2]] + Yη = [Y[1], -Y[2]] + ef = convert(ExponentialFamilyDistribution, M, p) + fisher_info = fisherinformation(ef) + return dot(Xη, fisher_info, 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 \ No newline at end of file diff --git a/src/with_metric.jl b/src/with_metric.jl index e033c84..1e0bf7f 100644 --- a/src/with_metric.jl +++ b/src/with_metric.jl @@ -55,6 +55,7 @@ function ManifoldsBase.norm(M::WithMetric{F,T,NaturalMetric}, p, X) where {F,T} return sqrt(inner(M, p, X, X)) end +# Default implementation for other distributions function ManifoldsBase.inner(M::WithMetric{F,T,NaturalMetric}, p, X, Y) where {F,T} ef = convert(ExponentialFamilyDistribution, M, p) fisher_info = fisherinformation(ef) diff --git a/test/metric_manifolds/mle_metric_manifolds_setuptests.jl b/test/metric_manifolds/mle_metric_manifolds_setuptests.jl index e6f52d1..8bc241c 100644 --- a/test/metric_manifolds/mle_metric_manifolds_setuptests.jl +++ b/test/metric_manifolds/mle_metric_manifolds_setuptests.jl @@ -16,6 +16,8 @@ function test_mle_works( ndistributions=10, backend_type=AutoForwardDiff(), kl_friendly=true, + use_sectional_curvature=false, + sectional_curvature_bound=0.0 ) rng = StableRNG(seed) @@ -39,8 +41,18 @@ function test_mle_works( return ManifoldDiff.gradient(M, (p) -> cost(M, p), p, backend) end - stepsize = DistanceOverGradients() - p_mle = gradient_descent(M, cost, grad, rand(rng, M); stepsize=stepsize) + 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 diff --git a/test/metric_manifolds/normal_tests.jl b/test/metric_manifolds/normal_tests.jl index 815913c..3325ad0 100644 --- a/test/metric_manifolds/normal_tests.jl +++ b/test/metric_manifolds/normal_tests.jl @@ -1,8 +1,8 @@ @testitem "Check `Normal` natural manifold" begin include("metric_manifolds_setuptests.jl") - test_metric_manifold(tol=1e-4, maximal_norm=0.3 ) do rng - return NormalMeanVariance(10randn(rng), rand(rng)+1) + test_metric_manifold(tol=1e-4) do rng + return NormalMeanVariance(10randn(rng), 10rand(rng)) end end From 0d3271ec16db5adcf50d90f3eb23d93b6058a689 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 19 Nov 2025 11:25:05 +0100 Subject: [PATCH 17/28] docs: add intention for WithMetric --- src/with_metric.jl | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/with_metric.jl b/src/with_metric.jl index 1e0bf7f..3b3ffd7 100644 --- a/src/with_metric.jl +++ b/src/with_metric.jl @@ -4,6 +4,16 @@ 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 From dc262407055b1794baa754a5a3010388c72eecfd Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 19 Nov 2025 13:15:48 +0100 Subject: [PATCH 18/28] fix: correct formula for inner --- src/with_metric.jl | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/with_metric.jl b/src/with_metric.jl index 3b3ffd7..200b943 100644 --- a/src/with_metric.jl +++ b/src/with_metric.jl @@ -65,9 +65,12 @@ function ManifoldsBase.norm(M::WithMetric{F,T,NaturalMetric}, p, X) where {F,T} return sqrt(inner(M, p, X, X)) end -# Default implementation for other distributions -function ManifoldsBase.inner(M::WithMetric{F,T,NaturalMetric}, p, X, Y) where {F,T} - ef = convert(ExponentialFamilyDistribution, M, p) +function ManifoldsBase.inner(M::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, Y) where {F} + # work on the underlying natural-parameters manifold when taking Jacobians + 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 + return dot(Xη, fisher_info, Yη) +end \ No newline at end of file From 9b07a5dd5120798650192bb93a2d9ff745985bd9 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 19 Nov 2025 13:16:39 +0100 Subject: [PATCH 19/28] fix: implement jacobian for normals --- src/ExponentialFamilyManifolds.jl | 1 + src/metric_manifolds/jacobians.jl | 65 +++++++++++++++++++++++++++++++ src/metric_manifolds/normal.jl | 19 +++------ 3 files changed, 72 insertions(+), 13 deletions(-) create mode 100644 src/metric_manifolds/jacobians.jl diff --git a/src/ExponentialFamilyManifolds.jl b/src/ExponentialFamilyManifolds.jl index ffd8731..63d426c 100644 --- a/src/ExponentialFamilyManifolds.jl +++ b/src/ExponentialFamilyManifolds.jl @@ -27,6 +27,7 @@ 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/jacobians.jl b/src/metric_manifolds/jacobians.jl new file mode 100644 index 0000000..540ad1d --- /dev/null +++ b/src/metric_manifolds/jacobians.jl @@ -0,0 +1,65 @@ +""" + 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_manifold_to_nat!(M, X_nat, X_manifold) + +Transform tangent vector from manifold coordinates to natural parameter coordinates. +""" +function jacobian_nat_to_manifold( + 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_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 \ No newline at end of file diff --git a/src/metric_manifolds/normal.jl b/src/metric_manifolds/normal.jl index 44a0be2..0ce2571 100644 --- a/src/metric_manifolds/normal.jl +++ b/src/metric_manifolds/normal.jl @@ -425,19 +425,12 @@ function ManifoldsBase.vector_transport_to!( return ManifoldsBase.parallel_transport_to!(M, Y, p, X, q) end -function ManifoldsBase.inner(M::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, Y) where {F} - # For NormalMeanVariance, manifold coords are (η₁, λ) but Fisher matrix is in (η₁, η₂) - # Transform tangent vectors: (dη₁, dλ) -> (dη₁, dη₂) where dη₂ = -dλ - Xη = [X[1], -X[2]] - Yη = [Y[1], -Y[2]] - ef = convert(ExponentialFamilyDistribution, M, p) - fisher_info = fisherinformation(ef) - return dot(Xη, fisher_info, 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) +""" +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 \ No newline at end of file From f9e3c8e2ca33b5afcdb51af5cac9fe82bce86c95 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 19 Nov 2025 13:48:55 +0100 Subject: [PATCH 20/28] style: :art: --- src/ExponentialFamilyManifolds.jl | 3 +- src/metric_manifolds/jacobians.jl | 16 +- src/metric_manifolds/normal.jl | 191 +++++++----------- src/with_metric.jl | 7 +- .../metric_manifolds_setuptests.jl | 7 +- .../mle_metric_manifolds_setuptests.jl | 13 +- test/metric_manifolds/normal_tests.jl | 2 +- 7 files changed, 98 insertions(+), 141 deletions(-) diff --git a/src/ExponentialFamilyManifolds.jl b/src/ExponentialFamilyManifolds.jl index 63d426c..7aa1403 100644 --- a/src/ExponentialFamilyManifolds.jl +++ b/src/ExponentialFamilyManifolds.jl @@ -1,6 +1,7 @@ module ExponentialFamilyManifolds -using BayesBase, ExponentialFamily, ManifoldsBase, Manifolds, Random, LinearAlgebra, StaticArrays +using BayesBase, + ExponentialFamily, ManifoldsBase, Manifolds, Random, LinearAlgebra, StaticArrays include("single_point_manifold.jl") include("natural_manifolds.jl") diff --git a/src/metric_manifolds/jacobians.jl b/src/metric_manifolds/jacobians.jl index 540ad1d..28d0588 100644 --- a/src/metric_manifolds/jacobians.jl +++ b/src/metric_manifolds/jacobians.jl @@ -4,8 +4,7 @@ Transform tangent vector from manifold coordinates to natural parameter coordinates. """ function jacobian_manifold_to_nat( - M::ExponentialFamilyManifolds.NaturalParametersManifold, - X_manifold, + M::ExponentialFamilyManifolds.NaturalParametersManifold, X_manifold ) X_nat = jacobian_manifold_to_nat!(M, zero_vector(M, X_manifold), X_manifold) return X_nat @@ -17,14 +16,12 @@ end Transform tangent vector from manifold coordinates to natural parameter coordinates. """ function jacobian_nat_to_manifold( - M::ExponentialFamilyManifolds.NaturalParametersManifold, - X_manifold, + 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_manifold, X_nat) @@ -33,8 +30,7 @@ For NormalMeanVariance: (dη₁, dη₂) → (dη₁, dλ) where dλ = -dη₂. """ function jacobian_nat_to_manifold!( ::ExponentialFamilyManifolds.NaturalParametersManifold{ - F, - ExponentialFamily.NormalMeanVariance, + F,ExponentialFamily.NormalMeanVariance }, X_manifold, X_nat, @@ -44,7 +40,6 @@ function jacobian_nat_to_manifold!( return X_manifold end - """ jacobian_manifold_to_nat!(M, X_nat, X_manifold) @@ -53,8 +48,7 @@ For NormalMeanVariance: (dη₁, dλ) → (dη₁, dη₂) where dη₂ = -dλ. """ function jacobian_manifold_to_nat!( ::ExponentialFamilyManifolds.NaturalParametersManifold{ - F, - ExponentialFamily.NormalMeanVariance, + F,ExponentialFamily.NormalMeanVariance }, X_nat, X_manifold, @@ -62,4 +56,4 @@ function jacobian_manifold_to_nat!( X_nat[1:1] .= X_manifold[1] X_nat[2:2] .= -X_manifold[2] return X_nat -end \ No newline at end of file +end diff --git a/src/metric_manifolds/normal.jl b/src/metric_manifolds/normal.jl index 0ce2571..432f23a 100644 --- a/src/metric_manifolds/normal.jl +++ b/src/metric_manifolds/normal.jl @@ -21,7 +21,7 @@ using RecursiveArrayTools: ArrayPartition # (η₁, η₂) -> (μ, σ²) @inline function _eta_to_meanvariance(η1, η2) - μ = -η1 / (2 * η2) + μ = -η1 / (2 * η2) σ² = -inv(2 * η2) return μ, σ² end @@ -37,19 +37,19 @@ end # here x = μ/√2, y = σ @inline function _eta_to_halfplane(η::SVector{2,T}) where {T} η1, η2 = η - μ, σ² = _eta_to_meanvariance(η1, η2) + μ, σ² = _eta_to_meanvariance(η1, η2) @assert σ² > zero(T) "NormalMeanVariance natural parameters must satisfy η₂ < 0" - σ = sqrt(σ²) - x = μ / sqrt(T(2)) - y = σ + σ = 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 + μ = sqrt(T(2)) * x + σ² = y^2 η1, η2 = _meanvariance_to_eta(μ, σ²) return SA[η1, η2] end @@ -61,15 +61,17 @@ end # [∂y/∂η₁ ∂y/∂η₂] @inline function _J_eta_to_xy(η::SVector{2}) η1, η2 = η - μ, σ² = _eta_to_meanvariance(η1, η2) - σ = sqrt(σ²) - c = sqrt(2.0) + μ, σ² = _eta_to_meanvariance(η1, η2) + σ = sqrt(σ²) + c = sqrt(2.0) a11 = -c / (4 * η2) - a12 = c * η1 / (4 * η2^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] + return SA[ + a11 a12; + a21 a22 + ] end # J_(x,y)→η evaluated at (x,y) @@ -77,15 +79,17 @@ end # [∂η₂/∂x ∂η₂/∂y] @inline function _J_xy_to_eta(xy::SVector{2}) x, y = xy - c = sqrt(2.0) + c = sqrt(2.0) y2 = y^2 y3 = y2 * y - b11 = c / y2 + b11 = c / y2 b12 = -2c * x / y3 b21 = 0.0 - b22 = 1.0 / y3 - return SA[b11 b12; - b21 b22] + b22 = 1.0 / y3 + return SA[ + b11 b12; + b21 b22 + ] end # --- Half-plane ↔ hyperboloid (curvature -1) --- @@ -117,31 +121,30 @@ end 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) + 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 + denom = u0 - u2 denom2 = denom^2 return SA[ - -u1/denom2 1/denom u1/denom2; - -1/denom2 0.0 1/denom2 + -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 +@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 = _mdot(W, W) m2 = max(m2, zero(T)) speed = sqrt(m2) if speed < eps(T) @@ -158,11 +161,11 @@ 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) + 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 @@ -172,12 +175,10 @@ 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}, + ν::SVector{3,T}, μ::SVector{3,T}, v::SVector{3,T} ) where {T<:Real} α = -_mdot(ν, μ) - num = _mdot(μ - α*ν, v) + num = _mdot(μ - α*ν, v) coeff = num / (α + one(T)) return v + coeff * (ν + μ) end @@ -187,39 +188,31 @@ end # ============================================================================ # Geodesic in η-coordinates (internal helper) -function geodesic_exact_eta( - t::Real, - η0::SVector{2,T}, - Xη0::SVector{2,T}, -) where {T<:Real} +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) + 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 + V0 = Jη_xy * Xη0 Jxy_u = _J_hp_to_hyp(xy0) - W0 = (Jxy_u * V0) / sqrt(T(2)) + W0 = (Jxy_u * V0) / sqrt(T(2)) # geodesic on hyperboloid - u_t = _hyperbolic_exp(u0, W0, T(t)) + u_t = _hyperbolic_exp(u0, W0, T(t)) # back to η xy_t = _hyp_to_hp(u_t) - η_t = _halfplane_to_eta(xy_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) +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) @@ -231,18 +224,18 @@ 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) + up = _hp_to_hyp(xyp) + uq = _hp_to_hyp(xyq) # log on hyperboloid - ξ = _hyperbolic_log(up, uq) + ξ = _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 * ξ + Vp = Ju_xy * ξ Jxy_η = _J_xy_to_eta(xyp) - Xη = (Jxy_η * Vp) * sqrt(T(2)) + Xη = (Jxy_η * Vp) * sqrt(T(2)) return Xη end @@ -252,7 +245,7 @@ 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η) + X = _eta_to_X(Xη) return X end @@ -263,10 +256,7 @@ end # Exponential map implementations function ManifoldsBase.exp!( - ::WithMetric{F,NormalMeanVariance,NaturalMetric}, - q, - p, - X, + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, q, p, X ) where {F} p_t = geodesic_exact(1, p, X) q[1] = p_t[1] @@ -274,21 +264,13 @@ function ManifoldsBase.exp!( return q end -function ManifoldsBase.exp( - ::WithMetric{F,NormalMeanVariance,NaturalMetric}, - p, - X, -) where {F} +function ManifoldsBase.exp(::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X) where {F} p_t = geodesic_exact(1, p, X) return ArrayPartition([p_t[1]], [p_t[2]]) end function ManifoldsBase.exp_fused!( - ::WithMetric{F,NormalMeanVariance,NaturalMetric}, - q, - p, - X, - t::Number, + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, q, p, X, t::Number ) where {F} p_t = geodesic_exact(t, p, X) q[1] = p_t[1] @@ -297,10 +279,7 @@ function ManifoldsBase.exp_fused!( end function ManifoldsBase.exp_fused( - ::WithMetric{F,NormalMeanVariance,NaturalMetric}, - p, - X, - t::Number, + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, t::Number ) where {F} p_t = geodesic_exact(t, p, X) return ArrayPartition([p_t[1]], [p_t[2]]) @@ -309,20 +288,13 @@ end # Log map implementations function ManifoldsBase.log!( - ::WithMetric{F,NormalMeanVariance,NaturalMetric}, - X, - p, - q, + ::WithMetric{F,NormalMeanVariance,NaturalMetric}, X, p, q ) where {F} X .= log_map(p, q) return X end -function ManifoldsBase.log( - ::WithMetric{F,NormalMeanVariance,NaturalMetric}, - p, - q, -) where {F} +function ManifoldsBase.log(::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, q) where {F} X = log_map(p, q) return ArrayPartition([X[1]], [X[2]]) end @@ -330,11 +302,7 @@ end # Parallel transport in natural parameters function ManifoldsBase.parallel_transport_to( - M::WithMetric{F,NormalMeanVariance,NaturalMetric}, - p, - X, - q; - kwargs..., + M::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, q; kwargs... ) where {F} ηp = _point_to_eta(p) ηq = _point_to_eta(q) @@ -343,35 +311,30 @@ function ManifoldsBase.parallel_transport_to( # p, q to hyperboloid xyp = _eta_to_halfplane(ηp) xyq = _eta_to_halfplane(ηq) - up = _hp_to_hyp(xyp) - uq = _hp_to_hyp(xyq) + 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η + Vp = Jη_xy * Xη Jxy_u = _J_hp_to_hyp(xyp) - Wp = (Jxy_u * Vp) / sqrt(2.0) + Wp = (Jxy_u * Vp) / sqrt(2.0) # parallel transport on hyperboloid - Wq = _hyperbolic_parallel_transport(up, uq, Wp) + 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 + Vq = Ju_xy * Wq Jxy_η = _J_xy_to_eta(xyq) - Xηq = (Jxy_η * Vq) * sqrt(2.0) + Xηq = (Jxy_η * Vq) * sqrt(2.0) - Xq = _eta_to_X(Xηq) + 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..., + M::WithMetric{F,NormalMeanVariance,NaturalMetric}, Y, p, X, q; kwargs... ) where {F} Y .= ManifoldsBase.parallel_transport_to(M, p, X, q; kwargs...) return Y @@ -380,23 +343,14 @@ end # Direction-based PT using exp function ManifoldsBase.parallel_transport_direction( - M::WithMetric{F,NormalMeanVariance,NaturalMetric}, - p, - X, - d; - kwargs..., + M::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, d; kwargs... ) where {F} q = ManifoldsBase.exp(M, p, d) return ManifoldsBase.parallel_transport_to(M, p, X, q; kwargs...) end function ManifoldsBase.parallel_transport_direction!( - M::WithMetric{F,NormalMeanVariance,NaturalMetric}, - Y, - p, - X, - d; - kwargs..., + M::WithMetric{F,NormalMeanVariance,NaturalMetric}, Y, p, X, d; kwargs... ) where {F} q = ManifoldsBase.exp(M, p, d) return ManifoldsBase.parallel_transport_to!(M, Y, p, X, q; kwargs...) @@ -430,7 +384,8 @@ 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} - +function ManifoldsBase.sectional_curvature_max( + ::WithMetric{F,NormalMeanVariance,NaturalMetric} +) where {F} return -0.5 -end \ No newline at end of file +end diff --git a/src/with_metric.jl b/src/with_metric.jl index 200b943..4bfd2a5 100644 --- a/src/with_metric.jl +++ b/src/with_metric.jl @@ -4,7 +4,6 @@ const _FWD = ManifoldsBase.SimpleForwardingType() abstract type MetricMode end struct NaturalMetric <: MetricMode end # Fisher info in active coordinates - """ WithMetric(::Type{T}, metric, dims, base, conditioner) @@ -65,7 +64,9 @@ 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,NormalMeanVariance,NaturalMetric}, p, X, Y) where {F} +function ManifoldsBase.inner( + M::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, Y +) where {F} # work on the underlying natural-parameters manifold when taking Jacobians natural_M = M.man Xη = jacobian_manifold_to_nat(natural_M, X) @@ -73,4 +74,4 @@ function ManifoldsBase.inner(M::WithMetric{F,NormalMeanVariance,NaturalMetric}, ef = convert(ExponentialFamilyDistribution, M.man, p) fisher_info = fisherinformation(ef) return dot(Xη, fisher_info, Yη) -end \ No newline at end of file +end diff --git a/test/metric_manifolds/metric_manifolds_setuptests.jl b/test/metric_manifolds/metric_manifolds_setuptests.jl index cdc3d2a..74ddc75 100644 --- a/test/metric_manifolds/metric_manifolds_setuptests.jl +++ b/test/metric_manifolds/metric_manifolds_setuptests.jl @@ -1,10 +1,11 @@ using StableRNGs, ExponentialFamily, ManifoldsBase, LinearAlgebra using Distributions -import ExponentialFamilyManifolds: - get_fisher_manifold, partition_point +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) +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 _ diff --git a/test/metric_manifolds/mle_metric_manifolds_setuptests.jl b/test/metric_manifolds/mle_metric_manifolds_setuptests.jl index 8bc241c..17a37a5 100644 --- a/test/metric_manifolds/mle_metric_manifolds_setuptests.jl +++ b/test/metric_manifolds/mle_metric_manifolds_setuptests.jl @@ -17,7 +17,7 @@ function test_mle_works( backend_type=AutoForwardDiff(), kl_friendly=true, use_sectional_curvature=false, - sectional_curvature_bound=0.0 + sectional_curvature_bound=0.0, ) rng = StableRNG(seed) @@ -42,16 +42,21 @@ function test_mle_works( end if use_sectional_curvature - stepsize = DistanceOverGradients(M; use_curvature=true, sectional_curvature_bound=sectional_curvature_max(M)) + 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); + M, + cost, + grad, + rand(rng, M); stepsize=stepsize, - stopping_criterion=stopping_criterion + stopping_criterion=stopping_criterion, ) ef_mle = convert(ExponentialFamilyDistribution, M, p_mle) if kl_friendly diff --git a/test/metric_manifolds/normal_tests.jl b/test/metric_manifolds/normal_tests.jl index 3325ad0..a6886ed 100644 --- a/test/metric_manifolds/normal_tests.jl +++ b/test/metric_manifolds/normal_tests.jl @@ -11,4 +11,4 @@ end test_mle_works() do rng return NormalMeanVariance(10randn(rng), 10rand(rng)) end -end \ No newline at end of file +end From c611d38a71c439a79870e695c96722889d30f4bc Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 19 Nov 2025 14:43:05 +0100 Subject: [PATCH 21/28] test: add check_vector_transport --- src/metric_manifolds/normal.jl | 3 +-- src/with_metric.jl | 1 - test/metric_manifolds/metric_manifolds_setuptests.jl | 3 +++ 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/metric_manifolds/normal.jl b/src/metric_manifolds/normal.jl index 432f23a..6d0e551 100644 --- a/src/metric_manifolds/normal.jl +++ b/src/metric_manifolds/normal.jl @@ -163,7 +163,7 @@ function _hyperbolic_log(u::SVector{3,T}, v::SVector{3,T}) where {T<:Real} end ip = _mdot(u, v) α = -ip - d = acosh(α) + d = acosh(α+1e-3) w = v - α*u # v - α u, α = -⟨u,v⟩ n2 = _mdot(w, w) if n2 <= zero(T) @@ -300,7 +300,6 @@ function ManifoldsBase.log(::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, end # Parallel transport in natural parameters - function ManifoldsBase.parallel_transport_to( M::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, q; kwargs... ) where {F} diff --git a/src/with_metric.jl b/src/with_metric.jl index 4bfd2a5..8b1519f 100644 --- a/src/with_metric.jl +++ b/src/with_metric.jl @@ -67,7 +67,6 @@ end function ManifoldsBase.inner( M::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, Y ) where {F} - # work on the underlying natural-parameters manifold when taking Jacobians natural_M = M.man Xη = jacobian_manifold_to_nat(natural_M, X) Yη = jacobian_manifold_to_nat(natural_M, Y) diff --git a/test/metric_manifolds/metric_manifolds_setuptests.jl b/test/metric_manifolds/metric_manifolds_setuptests.jl index 74ddc75..fa910bc 100644 --- a/test/metric_manifolds/metric_manifolds_setuptests.jl +++ b/test/metric_manifolds/metric_manifolds_setuptests.jl @@ -37,7 +37,10 @@ function test_metric_manifold( 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 From 571d2d933826ffdcd5b9c5653acee8c497f36c03 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 19 Nov 2025 15:52:16 +0100 Subject: [PATCH 22/28] fix: stop forwarding for normal to reduce amount of non-trivial code --- src/metric_manifolds/normal.jl | 141 +++++++++++++++++++++++++++------ 1 file changed, 115 insertions(+), 26 deletions(-) diff --git a/src/metric_manifolds/normal.jl b/src/metric_manifolds/normal.jl index 6d0e551..ae16872 100644 --- a/src/metric_manifolds/normal.jl +++ b/src/metric_manifolds/normal.jl @@ -1,10 +1,113 @@ 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 @@ -17,16 +120,12 @@ using RecursiveArrayTools: ArrayPartition @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]] -# --- Basic conversions: natural ↔ mean/variance ↔ half-plane --- - -# (η₁, η₂) -> (μ, σ²) @inline function _eta_to_meanvariance(η1, η2) μ = -η1 / (2 * η2) σ² = -inv(2 * η2) return μ, σ² end -# (μ, σ²) -> (η₁, η₂) @inline function _meanvariance_to_eta(μ, σ²) η1 = μ / σ² η2 = -inv(2 * σ²) @@ -163,7 +262,7 @@ function _hyperbolic_log(u::SVector{3,T}, v::SVector{3,T}) where {T<:Real} end ip = _mdot(u, v) α = -ip - d = acosh(α+1e-3) + d = acosh(α) w = v - α*u # v - α u, α = -⟨u,v⟩ n2 = _mdot(w, w) if n2 <= zero(T) @@ -249,12 +348,6 @@ function log_map(p, q) return X end -# ============================================================================ -# ManifoldsBase method overrides -# ============================================================================ - -# Exponential map implementations - function ManifoldsBase.exp!( ::WithMetric{F,NormalMeanVariance,NaturalMetric}, q, p, X ) where {F} @@ -264,10 +357,6 @@ function ManifoldsBase.exp!( return q end -function ManifoldsBase.exp(::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X) where {F} - p_t = geodesic_exact(1, p, X) - return ArrayPartition([p_t[1]], [p_t[2]]) -end function ManifoldsBase.exp_fused!( ::WithMetric{F,NormalMeanVariance,NaturalMetric}, q, p, X, t::Number @@ -278,12 +367,12 @@ function ManifoldsBase.exp_fused!( return q end -function ManifoldsBase.exp_fused( - ::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, t::Number -) where {F} - p_t = geodesic_exact(t, p, X) - return ArrayPartition([p_t[1]], [p_t[2]]) -end +# function ManifoldsBase.exp_fused( +# ::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, t::Number +# ) where {F} +# p_t = geodesic_exact(t, p, X) +# return ArrayPartition([p_t[1]], [p_t[2]]) +# end # Log map implementations @@ -294,10 +383,10 @@ function ManifoldsBase.log!( return X end -function ManifoldsBase.log(::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, q) where {F} - X = log_map(p, q) - return ArrayPartition([X[1]], [X[2]]) -end +# function ManifoldsBase.log(::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, q) where {F} +# X = log_map(p, q) +# return ArrayPartition([X[1]], [X[2]]) +# end # Parallel transport in natural parameters function ManifoldsBase.parallel_transport_to( From 4d32f52b2ffa8e511b320a27cba6bedfd182f31e Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 19 Nov 2025 16:01:19 +0100 Subject: [PATCH 23/28] fix: delete wrong forwarding implementation --- src/with_metric.jl | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/with_metric.jl b/src/with_metric.jl index 8b1519f..29ece4c 100644 --- a/src/with_metric.jl +++ b/src/with_metric.jl @@ -31,12 +31,6 @@ function with_natural_metric(M::NaturalParametersManifold{F}) where {F} 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}, ::Any, ::typeof(inner) -) where {F,T} = _STOP -@inline ManifoldsBase.get_forwarding_type( - ::WithMetric{F,T,NaturalMetric}, ::Any, ::typeof(norm) -) where {F,T} = _STOP @inline ManifoldsBase.get_forwarding_type( ::WithMetric{F,T,NaturalMetric}, ::typeof(inner) ) where {F,T} = _STOP From ce5241b6d7f5433da6156ba6e821f5cd6f94e869 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 19 Nov 2025 16:11:46 +0100 Subject: [PATCH 24/28] style: clean code --- src/metric_manifolds/normal.jl | 72 +++------------------------------- 1 file changed, 6 insertions(+), 66 deletions(-) diff --git a/src/metric_manifolds/normal.jl b/src/metric_manifolds/normal.jl index ae16872..d1f65f6 100644 --- a/src/metric_manifolds/normal.jl +++ b/src/metric_manifolds/normal.jl @@ -11,10 +11,14 @@ using RecursiveArrayTools: ArrayPartition ::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 + ::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 + ::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 @@ -348,16 +352,6 @@ function log_map(p, q) return X end -function ManifoldsBase.exp!( - ::WithMetric{F,NormalMeanVariance,NaturalMetric}, q, p, X -) where {F} - p_t = geodesic_exact(1, p, X) - q[1] = p_t[1] - q[2] = p_t[2] - return q -end - - function ManifoldsBase.exp_fused!( ::WithMetric{F,NormalMeanVariance,NaturalMetric}, q, p, X, t::Number ) where {F} @@ -367,15 +361,6 @@ function ManifoldsBase.exp_fused!( return q end -# function ManifoldsBase.exp_fused( -# ::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, t::Number -# ) where {F} -# p_t = geodesic_exact(t, p, X) -# return ArrayPartition([p_t[1]], [p_t[2]]) -# end - -# Log map implementations - function ManifoldsBase.log!( ::WithMetric{F,NormalMeanVariance,NaturalMetric}, X, p, q ) where {F} @@ -383,12 +368,6 @@ function ManifoldsBase.log!( return X end -# function ManifoldsBase.log(::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, q) where {F} -# X = log_map(p, q) -# return ArrayPartition([X[1]], [X[2]]) -# end - -# Parallel transport in natural parameters function ManifoldsBase.parallel_transport_to( M::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, q; kwargs... ) where {F} @@ -428,45 +407,6 @@ function ManifoldsBase.parallel_transport_to!( return Y end -# Direction-based PT using exp - -function ManifoldsBase.parallel_transport_direction( - M::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, d; kwargs... -) where {F} - q = ManifoldsBase.exp(M, p, d) - return ManifoldsBase.parallel_transport_to(M, p, X, q; kwargs...) -end - -function ManifoldsBase.parallel_transport_direction!( - M::WithMetric{F,NormalMeanVariance,NaturalMetric}, Y, p, X, d; kwargs... -) where {F} - q = ManifoldsBase.exp(M, p, d) - return ManifoldsBase.parallel_transport_to!(M, Y, p, X, q; kwargs...) -end - -# vector_transport_to aliases ParallelTransport, like in your Bernoulli code - -function ManifoldsBase.vector_transport_to( - M::WithMetric{F,NormalMeanVariance,NaturalMetric}, - p, - X, - q, - ::ManifoldsBase.ParallelTransport, -) where {F} - return ManifoldsBase.parallel_transport_to(M, p, X, q) -end - -function ManifoldsBase.vector_transport_to!( - M::WithMetric{F,NormalMeanVariance,NaturalMetric}, - Y, - p, - X, - q, - ::ManifoldsBase.ParallelTransport, -) where {F} - return ManifoldsBase.parallel_transport_to!(M, Y, p, X, q) -end - """ Sectional curvature for Normal manifold The Normal Fisher manifold is isometric to hyperbolic space with curvature -1/2 From 6bb4b922dbcccadb4d3cc67832b10d17cffc12c4 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 19 Nov 2025 16:22:26 +0100 Subject: [PATCH 25/28] fix: clean bernoulli --- src/metric_manifolds/bernoulli.jl | 156 +++++++++++++++++++----------- 1 file changed, 97 insertions(+), 59 deletions(-) diff --git a/src/metric_manifolds/bernoulli.jl b/src/metric_manifolds/bernoulli.jl index 6c90821..e6cfc6c 100644 --- a/src/metric_manifolds/bernoulli.jl +++ b/src/metric_manifolds/bernoulli.jl @@ -1,3 +1,95 @@ +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 @@ -38,42 +130,19 @@ function _log_map_bernoulli(x0, x1) return v0 end -# Exponential map implementations - these override the default forwarding behavior -function ManifoldsBase.exp!(::WithMetric{F,Bernoulli,NaturalMetric}, q, p, X) where {F} - x, _ = _geodesic_exact_bernoulli(1, p[1], X[1]) - q .= x - return q -end - -function ManifoldsBase.exp(::WithMetric{F,Bernoulli,NaturalMetric}, p, X) where {F} - x, _ = _geodesic_exact_bernoulli(1, p[1], X[1]) - return [x] -end - function ManifoldsBase.exp_fused!( ::WithMetric{F,Bernoulli,NaturalMetric}, q, p, X, t::Number ) where {F} - x, _ = _geodesic_exact_bernoulli(1, p[1], X[1]) + x, _ = _geodesic_exact_bernoulli(t, p[1], X[1]) q .= x return q end -function ManifoldsBase.exp_fused( - ::WithMetric{F,Bernoulli,NaturalMetric}, p, X, t::Number -) where {F} - x, _ = _geodesic_exact_bernoulli(t, p[1], X[1]) - return [x] -end - function ManifoldsBase.log!(::WithMetric{F,Bernoulli,NaturalMetric}, X, p, q) where {F} X .= _log_map_bernoulli(p[1], q[1]) return X end -function ManifoldsBase.log(::WithMetric{F,Bernoulli,NaturalMetric}, p, q) where {F} - return _log_map_bernoulli(p[1], q[1]) -end - _sqrtg_eta(η) = 0.5 .* sech.(η ./ 2) function ManifoldsBase.parallel_transport_to( @@ -81,43 +150,12 @@ function ManifoldsBase.parallel_transport_to( ) where {F} s_p = _sqrtg_eta(p) s_q = _sqrtg_eta(q) - X .* (s_p ./ s_q) + return X .* (s_p ./ s_q) end function ManifoldsBase.parallel_transport_to!( - ::WithMetric{F,Bernoulli,NaturalMetric}, Y, p, X, q; kwargs... -) where {F} - # error("Call here") - s_p = _sqrtg_eta(p) - s_q = _sqrtg_eta(q) - @. Y = X * (s_p / s_q) - Y -end - -# Parallel transport in a direction -function ManifoldsBase.parallel_transport_direction( - M::WithMetric{F,Bernoulli,NaturalMetric}, p, X, d; kwargs... -) where {F} - q = exp(M, p, d) - return parallel_transport_to(M, p, X, q) -end - -function ManifoldsBase.parallel_transport_direction!( - M::WithMetric{F,Bernoulli,NaturalMetric}, Y, p, X, d; kwargs... -) where {F} - q = exp(M, p, d) - return parallel_transport_to!(M, Y, p, X, q) -end - -# Vector transport with ParallelTransport method - this is what check_geodesic uses -function ManifoldsBase.vector_transport_to( - M::WithMetric{F,Bernoulli,NaturalMetric}, p, X, q, ::ManifoldsBase.ParallelTransport -) where {F} - return parallel_transport_to(M, p, X, q) -end - -function ManifoldsBase.vector_transport_to!( - M::WithMetric{F,Bernoulli,NaturalMetric}, Y, p, X, q, ::ManifoldsBase.ParallelTransport + M::WithMetric{F,Bernoulli,NaturalMetric}, Y, p, X, q; kwargs... ) where {F} - return parallel_transport_to!(M, Y, p, X, q) + Y .= ManifoldsBase.parallel_transport_to(M, p, X, q; kwargs...) + return Y end From 39eb554c85951d3d47f39ed5545af789e7706222 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 19 Nov 2025 16:39:33 +0100 Subject: [PATCH 26/28] fix: check that jacob operators are inverses of each other --- .../metric_manifolds_setuptests.jl | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/test/metric_manifolds/metric_manifolds_setuptests.jl b/test/metric_manifolds/metric_manifolds_setuptests.jl index fa910bc..88fe47a 100644 --- a/test/metric_manifolds/metric_manifolds_setuptests.jl +++ b/test/metric_manifolds/metric_manifolds_setuptests.jl @@ -32,6 +32,28 @@ function test_metric_manifold( @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 + #respect fisher metric for _ in 1:test_points p = rand(rng, M) From f5cfa2b2391dc5db6b030ecd952830ff6abd866c Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 19 Nov 2025 16:40:31 +0100 Subject: [PATCH 27/28] docs: readiable jacobian_nat_to_manifold --- src/metric_manifolds/bernoulli.jl | 2 +- src/metric_manifolds/jacobians.jl | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/metric_manifolds/bernoulli.jl b/src/metric_manifolds/bernoulli.jl index e6cfc6c..f11ae3e 100644 --- a/src/metric_manifolds/bernoulli.jl +++ b/src/metric_manifolds/bernoulli.jl @@ -149,7 +149,7 @@ 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) + s_q = _sqrtg_eta(q) return X .* (s_p ./ s_q) end diff --git a/src/metric_manifolds/jacobians.jl b/src/metric_manifolds/jacobians.jl index 28d0588..77dedef 100644 --- a/src/metric_manifolds/jacobians.jl +++ b/src/metric_manifolds/jacobians.jl @@ -11,15 +11,15 @@ function jacobian_manifold_to_nat( end """ - jacobian_manifold_to_nat!(M, X_nat, X_manifold) + jacobian_nat_to_manifold(M, X_nat) -Transform tangent vector from manifold coordinates to natural parameter coordinates. +Transform tangent vector from natural parameter coordinates to manifold coordinates. """ function jacobian_nat_to_manifold( - M::ExponentialFamilyManifolds.NaturalParametersManifold, X_manifold + M::ExponentialFamilyManifolds.NaturalParametersManifold, X_nat ) - X_nat = jacobian_manifold_to_nat!(M, zero_vector(M, X_manifold), X_manifold) - return X_nat + X_manifold = jacobian_nat_to_manifold!(M, zero_vector(M, X_nat), X_nat) + return X_manifold end """ From 419ae11b7acecf1051d49c62a3ec5f414b754dc4 Mon Sep 17 00:00:00 2001 From: Mykola Lukashchuk Date: Wed, 19 Nov 2025 16:49:15 +0100 Subject: [PATCH 28/28] fix: correct metric on Bernoulli --- src/metric_manifolds/jacobians.jl | 34 +++++++++++++++++++ src/with_metric.jl | 4 +-- test/metric_manifolds/bernoulli_tests.jl | 2 +- .../metric_manifolds_setuptests.jl | 16 +++++++++ 4 files changed, 53 insertions(+), 3 deletions(-) diff --git a/src/metric_manifolds/jacobians.jl b/src/metric_manifolds/jacobians.jl index 77dedef..10b3572 100644 --- a/src/metric_manifolds/jacobians.jl +++ b/src/metric_manifolds/jacobians.jl @@ -57,3 +57,37 @@ function jacobian_manifold_to_nat!( 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/with_metric.jl b/src/with_metric.jl index 29ece4c..6c929f3 100644 --- a/src/with_metric.jl +++ b/src/with_metric.jl @@ -59,8 +59,8 @@ function ManifoldsBase.norm(M::WithMetric{F,T,NaturalMetric}, p, X) where {F,T} end function ManifoldsBase.inner( - M::WithMetric{F,NormalMeanVariance,NaturalMetric}, p, X, Y -) where {F} + 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) diff --git a/test/metric_manifolds/bernoulli_tests.jl b/test/metric_manifolds/bernoulli_tests.jl index 3ac51d4..3e83aef 100644 --- a/test/metric_manifolds/bernoulli_tests.jl +++ b/test/metric_manifolds/bernoulli_tests.jl @@ -1,6 +1,6 @@ @testitem "Check `Bernoulli` fisher manifold" begin include("metric_manifolds_setuptests.jl") - test_metric_manifold(maximal_norm=0.4) do rng + test_metric_manifold(maximal_norm=0.3) 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 index 88fe47a..1880888 100644 --- a/test/metric_manifolds/metric_manifolds_setuptests.jl +++ b/test/metric_manifolds/metric_manifolds_setuptests.jl @@ -53,6 +53,22 @@ function test_metric_manifold( ) @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