diff --git a/Project.toml b/Project.toml index b584cc3..efef5b7 100644 --- a/Project.toml +++ b/Project.toml @@ -6,6 +6,7 @@ version = "2.0.0" [deps] BayesBase = "b4ee3484-f114-42fe-b91c-797d54a0c67e" ExponentialFamily = "62312e5e-252a-4322-ace9-a5f4bf9b357b" +FastCholesky = "2d5283b6-8564-42b6-bb00-83ed8e915756" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e" ManifoldsBase = "3362f125-f0bb-47a3-aa74-596ffd7ef2fb" @@ -13,10 +14,16 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" +[weakdeps] +ManifoldDiff = "af67fdf4-a580-4b9f-bbec-742ef357defd" + [compat] -BayesBase = "1.3" +ADTypes = "1.14.0" +BayesBase = "1.5.4" ExponentialFamily = "2.0.0" +FastCholesky = "1.3.1" LinearAlgebra = "1.10" +ManifoldDiff = "0.4.2" Manifolds = "0.10" ManifoldsBase = "1" Random = "1.10" @@ -25,6 +32,7 @@ Static = "0.8, 1" julia = "1.10" [extras] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" @@ -36,4 +44,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Aqua", "Distributions", "JET", "Test", "ReTestItems", "StableRNGs", "StaticArrays", "Manopt", "ForwardDiff"] +test = ["ADTypes", "Aqua", "Distributions", "JET", "Test", "ReTestItems", "StableRNGs", "StaticArrays", "Manopt", "ForwardDiff"] diff --git a/src/natural_manifolds.jl b/src/natural_manifolds.jl index d64cf3d..29f5667 100644 --- a/src/natural_manifolds.jl +++ b/src/natural_manifolds.jl @@ -1,27 +1,119 @@ using ManifoldsBase, Manifolds, Static, RecursiveArrayTools, Random, ExponentialFamily +using FastCholesky import ExponentialFamily: exponential_family_typetag + +struct ChartNOrderRetraction{Order,E} <: AbstractRetractionMethod + extra::E +end + +function ChartNOrderRetraction{O}() where {O} + return ChartNOrderRetraction{O,Nothing}(nothing) +end + +const FirstOrderRetraction = ChartNOrderRetraction{1} +const SecondOrderRetraction = ChartNOrderRetraction{2} + +""" + SecondOrderRetraction(; backend=nothing) + +Create a second-order retraction method that uses Christoffel symbols to compute +a more accurate retraction. If a backend is provided, it will be used for any +automatic differentiation needed to compute the Christoffel symbols. + +# Arguments +- `backend`: Optional backend for automatic differentiation (e.g., `ADTypes.AutoForwardDiff()`) +""" +function SecondOrderRetraction(; backend=nothing) + return ChartNOrderRetraction{2,typeof(backend)}(backend) +end + +""" + FisherInformationMetric <: RiemannianMetric + +Specifier that we need to use the Fisher information metric. +""" +struct FisherInformationMetric{R} <: RiemannianMetric + default_retraction::R +end + +function FisherInformationMetric() + retraction = FirstOrderRetraction() + return FisherInformationMetric{typeof(retraction)}(retraction) +end + +""" + BaseMetric <: RiemannianMetric + +Specifier that we need to use the metric from the base manifold. +""" +struct BaseMetric <: RiemannianMetric end + +""" + getdefaultmetric(::Type{T}) where {T} + +Returns the default metric for the distribution of type `T`. +""" +function getdefaultmetric(::Type{T}) where {T} + return FisherInformationMetric() +end + """ NaturalParametersManifold(::Type{T}, dims, base, conditioner) The manifold for the natural parameters of the distribution of type `T` with dimensions `dims`. An internal structure, use `get_natural_manifold` to create an instance of a manifold for the natural parameters of distribution of type `T`. """ -struct NaturalParametersManifold{𝔽,T,D,M,C} <: AbstractDecoratorManifold{𝔽} +struct NaturalParametersManifold{𝔽,T,D,M,C,MT} <: AbstractDecoratorManifold{𝔽} dims::D base::M conditioner::C + metric::MT end getdims(M::NaturalParametersManifold) = M.dims getbase(M::NaturalParametersManifold) = M.base getconditioner(M::NaturalParametersManifold) = M.conditioner +getmetric(M::NaturalParametersManifold) = M.metric # The `NaturalParametersManifold` simply adds extra properties to the `base` and # acts as a "decorator" -@inline ManifoldsBase.active_traits(f::F, ::NaturalParametersManifold, ::Any...) where {F} = - ManifoldsBase.IsExplicitDecorator() +function select_skip_methods( + ::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,BaseMetric} +) where {F,𝔽,T,D,MB,C} + return ManifoldsBase.IsExplicitDecorator() +end + +function select_skip_methods( + f::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,<:FisherInformationMetric} +) where {F,𝔽,T,D,MB,C} + if f in ( + ManifoldsBase.retract, + ManifoldsBase.retract!, + ManifoldsBase.retract_fused, + ManifoldsBase.retract_fused!, + Manifolds.local_metric, + Manifolds.local_metric_jacobian, + Manifolds.inverse_local_metric, + Manifolds.default_retraction_method, + Manifolds.get_basis_default, + Manifolds.christoffel_symbols_second, + Manifolds.christoffel_symbols_first, + Manifolds.representation_size + ) + return ManifoldsBase.EmptyTrait() + else + return ManifoldsBase.IsExplicitDecorator() + end +end + +@inline function ManifoldsBase.active_traits( + f::F, M::NaturalParametersManifold, args... +) where {F} + return select_skip_methods(f, M) +end + @inline ManifoldsBase.decorated_manifold(M::NaturalParametersManifold) = M.base function ExponentialFamily.exponential_family_typetag( @@ -31,9 +123,15 @@ function ExponentialFamily.exponential_family_typetag( end function NaturalParametersManifold( - ::Type{T}, dims::D, base::M, conditioner::C=nothing -) where {T,𝔽,D,M<:AbstractManifold{𝔽},C} - return NaturalParametersManifold{𝔽,T,D,M,C}(dims, base, conditioner) + ::Type{T}, + dims::D, + base::M, + conditioner::C=nothing, + metric::MT=getdefaultmetric(T), +) where {T,𝔽,D,M<:AbstractManifold{𝔽},C,MT} + return NaturalParametersManifold{𝔽,T,D,M,C,MT}( + dims, base, conditioner, metric + ) end """ @@ -52,9 +150,11 @@ julia> ExponentialFamilyManifolds.get_natural_manifold(MvNormalMeanCovariance, ( true ``` """ -function get_natural_manifold(::Type{T}, dims, conditioner=nothing) where {T} +function get_natural_manifold( + ::Type{T}, dims, conditioner=nothing, metric=getdefaultmetric(T) +) where {T} return NaturalParametersManifold( - T, dims, get_natural_manifold_base(T, dims, conditioner), conditioner + T, dims, get_natural_manifold_base(T, dims, conditioner), conditioner, metric ) end @@ -88,3 +188,75 @@ function Base.convert( exponential_family_typetag(M), p, getconditioner(M), nothing ) end + + +function ManifoldsBase.default_retraction_method( + M::NaturalParametersManifold{𝔽,TD,D,BM,C,<:FisherInformationMetric}, ::Type{T} +) where {𝔽,T,TD,D,BM,C} + return getmetric(M).default_retraction +end + +function ManifoldsBase.retract_fused!( + ::NaturalParametersManifold, q, p, X, t::Number, method::FirstOrderRetraction +) + q .= p .+ t .* X + return q +end + +function ManifoldsBase.retract!( + M::NaturalParametersManifold, q, p, X, method::FirstOrderRetraction +) + return ManifoldsBase.retract_fused!(M, q, p, X, one(eltype(X)), method) +end + +function ManifoldsBase.retract_fused!( + M::NaturalParametersManifold{𝔽,T,D,BM,C,<:FisherInformationMetric}, + q, + p, + X, + t::Number, + method::SecondOrderRetraction, +) where {𝔽,T,D,BM,C} + basis = ManifoldsBase.get_basis_default(M, p) + Γ = Manifolds.christoffel_symbols_second(M, p, basis; backend=method.extra) + + Δ = similar(p) + Manifolds.@einsum Δ[k] = -0.5 * Γ[k, i, j] * (t * X[i]) * (t * X[j]) + q .= p .+ t .* X .+ Δ + return q +end + +function ManifoldsBase.retract!( + M::NaturalParametersManifold, q, p, X, method::SecondOrderRetraction +) + return ManifoldsBase.retract_fused!(M, q, p, X, one(eltype(X)), method) +end + +struct NaturalBasis{𝔽,VST<:VectorSpaceType} <: AbstractBasis{𝔽,VST} + vector_space::VST +end + +NaturalBasis(𝔽=ℝ, vs::VectorSpaceType=TangentSpaceType()) = NaturalBasis{𝔽,typeof(vs)}(vs) +function NaturalBasis{𝔽}(vs::VectorSpaceType=TangentSpaceType()) where {𝔽} + return NaturalBasis{𝔽,typeof(vs)}(vs) +end + +function ManifoldsBase.get_basis_default( + ::NaturalParametersManifold{𝔽,T,D,MB,C,<:FisherInformationMetric}, p +) where {𝔽,T,D,MB,C} + return NaturalBasis{𝔽}() +end + +function Manifolds.local_metric( + M::NaturalParametersManifold{𝔽,T,D,MB,C,<:FisherInformationMetric}, p, ::NaturalBasis +) where {𝔽,T,D,MB,C} + ef = convert(ExponentialFamilyDistribution, M, p) + return ExponentialFamily.fisherinformation(ef) +end + +function Manifolds.inverse_local_metric( + M::NaturalParametersManifold{𝔽,T,D,MB,C,<:FisherInformationMetric}, p, ::NaturalBasis +) where {𝔽,T,D,MB,C} + ef = convert(ExponentialFamilyDistribution, M, p) + return cholinv(ExponentialFamily.fisherinformation(ef)) +end diff --git a/src/natural_manifolds/beta.jl b/src/natural_manifolds/beta.jl index a9849f6..df949f8 100644 --- a/src/natural_manifolds/beta.jl +++ b/src/natural_manifolds/beta.jl @@ -17,3 +17,5 @@ Converts the `point` to a compatible representation for the natural manifold of function partition_point(::Type{Beta}, ::Tuple{}, p, conditioner=nothing) return ArrayPartition(view(p, 1:1), view(p, 2:2)) end + +Manifolds.representation_size(::NaturalParametersManifold{𝔽, Beta}) where {𝔽} = (2,) \ No newline at end of file diff --git a/src/natural_manifolds/categorical.jl b/src/natural_manifolds/categorical.jl index f35b8e2..5de7736 100644 --- a/src/natural_manifolds/categorical.jl +++ b/src/natural_manifolds/categorical.jl @@ -16,3 +16,7 @@ Converts the `point` to a compatible representation for the natural manifold of function partition_point(::Type{Categorical}, ::Tuple{}, p, conditioner=nothing) return ArrayPartition(view(p, 1:(conditioner - 1)), view(p, conditioner:conditioner)) end + +function Manifolds.representation_size(M::NaturalParametersManifold{𝔽, Categorical}) where {𝔽} + return (getconditioner(M),) +end diff --git a/src/natural_manifolds/dirichlet.jl b/src/natural_manifolds/dirichlet.jl index dd0e230..58f2f96 100644 --- a/src/natural_manifolds/dirichlet.jl +++ b/src/natural_manifolds/dirichlet.jl @@ -19,3 +19,8 @@ function partition_point(::Type{Dirichlet}, dims::Tuple{Int}, p, conditioner=not # See comment in `get_natural_manifold_base` for `Dirichlet` return ArrayPartition(p') end + +function Manifolds.representation_size(M::NaturalParametersManifold{𝔽, Dirichlet}) where {𝔽} + dims = getdims(M) + return (first(dims),) +end \ No newline at end of file diff --git a/src/natural_manifolds/gamma.jl b/src/natural_manifolds/gamma.jl index d910c70..412f3bc 100644 --- a/src/natural_manifolds/gamma.jl +++ b/src/natural_manifolds/gamma.jl @@ -18,3 +18,5 @@ Converts the `point` to a compatible representation for the natural manifold of function partition_point(::Type{Gamma}, ::Tuple{}, p, conditioner=nothing) return ArrayPartition(view(p, 1:1), view(p, 2:2)) end + +Manifolds.representation_size(::NaturalParametersManifold{𝔽, Gamma}) where {𝔽} = (2,) \ No newline at end of file diff --git a/src/natural_manifolds/inverse_gamma.jl b/src/natural_manifolds/inverse_gamma.jl index d046e2b..57734dd 100644 --- a/src/natural_manifolds/inverse_gamma.jl +++ b/src/natural_manifolds/inverse_gamma.jl @@ -21,3 +21,5 @@ function partition_point( ) return ArrayPartition(view(p, 1:1), view(p, 2:2)) end + +Manifolds.representation_size(::NaturalParametersManifold{𝔽, ExponentialFamily.GammaInverse}) where {𝔽} = (2,) \ No newline at end of file diff --git a/src/natural_manifolds/lognormal.jl b/src/natural_manifolds/lognormal.jl index 90c7104..67e6532 100644 --- a/src/natural_manifolds/lognormal.jl +++ b/src/natural_manifolds/lognormal.jl @@ -15,3 +15,5 @@ Converts the `point` to a compatible representation for the natural manifold of function partition_point(::Type{LogNormal}, ::Tuple{}, p, conditioner=nothing) return ArrayPartition(view(p, 1:1), view(p, 2:2)) end + +Manifolds.representation_size(::NaturalParametersManifold{𝔽, LogNormal}) where {𝔽} = (2,) \ No newline at end of file diff --git a/src/natural_manifolds/normal.jl b/src/natural_manifolds/normal.jl index 9e7b40f..d0a0321 100644 --- a/src/natural_manifolds/normal.jl +++ b/src/natural_manifolds/normal.jl @@ -66,3 +66,13 @@ function partition_point( k = first(dims) return ArrayPartition(view(p, 1:k), view(p, (k + 1):(k + 1))) end + +function getdefaultmetric(::Type{MvNormalMeanCovariance}) + return BaseMetric() +end + +Manifolds.representation_size(::NaturalParametersManifold{𝔽, NormalMeanVariance}) where {𝔽} = (2,) + +function Manifolds.representation_size(M::NaturalParametersManifold{𝔽, MvNormalMeanScalePrecision}) where {𝔽} + return (M.dims[1] +1, ) +end \ No newline at end of file diff --git a/src/natural_manifolds/wishart.jl b/src/natural_manifolds/wishart.jl index 0749200..5269d17 100644 --- a/src/natural_manifolds/wishart.jl +++ b/src/natural_manifolds/wishart.jl @@ -21,3 +21,7 @@ function partition_point( k = first(dims) return ArrayPartition(view(p, 1:1), reshape(view(p, 2:(1 + k^2)), (k, k))) end + +function getdefaultmetric(::Type{ExponentialFamily.WishartFast}) + return BaseMetric() +end \ No newline at end of file diff --git a/test/natural_manifolds/bernoulli_tests.jl b/test/natural_manifolds/bernoulli_tests.jl index 82c7292..f382037 100644 --- a/test/natural_manifolds/bernoulli_tests.jl +++ b/test/natural_manifolds/bernoulli_tests.jl @@ -5,3 +5,25 @@ return Bernoulli(rand(rng)) end end + +@testitem "Check SecondOrderRetraction" begin + include("natural_manifolds_setuptests.jl") + + using ADTypes: AutoForwardDiff + + rng = StableRNG(42) + M = ExponentialFamilyManifolds.get_natural_manifold(Bernoulli, ()) + p = rand(rng, M) + X = rand(rng, M) + basis = ExponentialFamilyManifolds.NaturalBasis() + + @show Manifolds.local_metric(M, p, basis) + @show Manifolds.local_metric_jacobian(M, p, basis, backend=AutoForwardDiff()) + q = retract( + M, + p, + X, + ExponentialFamilyManifolds.SecondOrderRetraction(; backend=AutoForwardDiff()), + ) + @show q +end diff --git a/test/natural_manifolds/beta_tests.jl b/test/natural_manifolds/beta_tests.jl index 5df840e..7d2b68a 100644 --- a/test/natural_manifolds/beta_tests.jl +++ b/test/natural_manifolds/beta_tests.jl @@ -4,4 +4,5 @@ test_natural_manifold() do rng return Beta(10rand(rng), 10rand(rng)) end + end diff --git a/test/natural_manifolds/categorical_tests.jl b/test/natural_manifolds/categorical_tests.jl index dce5b2e..053bf61 100644 --- a/test/natural_manifolds/categorical_tests.jl +++ b/test/natural_manifolds/categorical_tests.jl @@ -1,7 +1,7 @@ @testitem "Check `Categorical` natural manifold" begin include("natural_manifolds_setuptests.jl") - test_natural_manifold() do rng + test_natural_manifold(test_injectivity_radius=false) do rng p = rand(rng, 10) normalize!(p, 1) return Categorical(p) diff --git a/test/natural_manifolds/natural_manifolds_setuptests.jl b/test/natural_manifolds/natural_manifolds_setuptests.jl index 6ebb980..de00742 100644 --- a/test/natural_manifolds/natural_manifolds_setuptests.jl +++ b/test/natural_manifolds/natural_manifolds_setuptests.jl @@ -1,9 +1,11 @@ -using StableRNGs, ExponentialFamily, ManifoldsBase, LinearAlgebra - +using StableRNGs, ExponentialFamily, Manifolds, ManifoldsBase, LinearAlgebra +using ADTypes: AutoForwardDiff import ExponentialFamilyManifolds: get_natural_manifold, partition_point -function test_natural_manifold(f; seed=42, ndistributions=100) +using FastCholesky + +function test_natural_manifold(f; seed=42, ndistributions=100, test_metric=true, test_injectivity_radius=true, fisher_fd_friendly=true) rng = StableRNG(seed) foreach(1:ndistributions) do _ @@ -17,5 +19,50 @@ function test_natural_manifold(f; seed=42, ndistributions=100) η = partition_point(T, dims, getnaturalparameters(ef), getconditioner(ef)) @test is_point(M, η, error=:error) + + if test_metric && M.metric isa ExponentialFamilyManifolds.FisherInformationMetric + @test ManifoldsBase.get_basis_default(M, η) isa + ExponentialFamilyManifolds.NaturalBasis + @test Manifolds.local_metric(M, η, ManifoldsBase.get_basis_default(M, η)) ≈ + ExponentialFamily.fisherinformation(ef) + @test Manifolds.inverse_local_metric(M, η, ManifoldsBase.get_basis_default(M, η)) ≈ + cholinv(ExponentialFamily.fisherinformation(ef)) + end + end + + distribution = f(rng) + sample = rand(rng, distribution) + dims = size(sample) + ef = convert(ExponentialFamilyDistribution, distribution) + + T = ExponentialFamily.exponential_family_typetag(ef) + M = get_natural_manifold(T, dims, getconditioner(ef)) + if test_metric && M.metric isa ExponentialFamilyManifolds.FisherInformationMetric + @testset "Testing $(typeof(M)) with retractions" begin + pts = [rand(rng, M) for _ in 1:5] + + if fisher_fd_friendly + retraction_methods = [ + ExponentialFamilyManifolds.FirstOrderRetraction(), + ExponentialFamilyManifolds.SecondOrderRetraction(backend=AutoForwardDiff()), + ] + else + retraction_methods = [ + ExponentialFamilyManifolds.FirstOrderRetraction(), + ] + end + + Manifolds.test_manifold( + M, + pts; + test_exp_log=true, + default_inverse_retraction_method=nothing, + test_default_vector_transport=false, + retraction_methods=retraction_methods, + inverse_retraction_methods=[], + test_injectivity_radius=test_injectivity_radius, + ) + end end end + \ No newline at end of file diff --git a/test/natural_manifolds/normal_tests.jl b/test/natural_manifolds/normal_tests.jl index 561fc99..caa5fa7 100644 --- a/test/natural_manifolds/normal_tests.jl +++ b/test/natural_manifolds/normal_tests.jl @@ -9,19 +9,27 @@ end @testitem "Check `MvNormal` natural manifold" begin include("natural_manifolds_setuptests.jl") - test_natural_manifold() do rng + test_natural_manifold(; test_metric=false) do rng k = rand(rng, 1:10) m = randn(k) L = LowerTriangular(randn(k, k)) C = L * L' + k * I return MvNormalMeanCovariance(m, C) end + + using ExponentialFamilyManifolds + using ExponentialFamily + + @test ExponentialFamilyManifolds.getdefaultmetric(MvNormalMeanCovariance) isa + ExponentialFamilyManifolds.BaseMetric + M = ExponentialFamilyManifolds.get_natural_manifold(MvNormalMeanCovariance, (3,)) + @test M.metric isa ExponentialFamilyManifolds.BaseMetric end @testitem "Check `MvNormalMeanScalePrecision` natural manifold" begin include("natural_manifolds_setuptests.jl") - test_natural_manifold() do rng + test_natural_manifold(fisher_fd_friendly=false) do rng k = rand(rng, 1:10) m = randn(rng, k) γ = rand(rng)^2 + 1 diff --git a/test/natural_manifolds_tests.jl b/test/natural_manifolds_tests.jl index 2342938..b230f71 100644 --- a/test/natural_manifolds_tests.jl +++ b/test/natural_manifolds_tests.jl @@ -76,3 +76,17 @@ end end end + +@testitem "Natural manifold properties: BaseMetric" begin + using Distributions, ExponentialFamily, Manifolds, ManifoldsBase, StableRNGs, JET, LinearAlgebra + + import ExponentialFamilyManifolds: BaseMetric + + M = ExponentialFamilyManifolds.get_natural_manifold(Beta, (), nothing, BaseMetric()) + @test M.metric isa BaseMetric + @test ExponentialFamilyManifolds.select_skip_methods(ManifoldsBase.retract, M) == ManifoldsBase.IsExplicitDecorator() + p = rand(StableRNG(42), M) + q = copy(p) + @test ManifoldsBase.retract!(M, q, p, p) ≈ ManifoldsBase.retract!(M.base, q, p, p) +end +