Merge branch 'main' of https://github.com/biaslab/LogExpectations

Author: Nimrais

src/ClosedFormExpectations.jl

@@ -55,7 +55,12 @@ include("Exponential/Exponential.jl")
 include("Normal/expectation.jl")
 include("Normal/williams/normal.jl")
 include("Normal/williams/normal_mean_variance.jl")
+include("Normal/williams/ef_parametrization.jl")
+
 # gamma
 include("Gamma/Gamma.jl")
 
+# exponetial family distribution interface
+include("exponential_family_interface.jl")
+
 end

src/Normal/williams/ef_parametrization.jl

@@ -0,0 +1,10 @@
+using StaticArrays
+import Distributions: var, mean, std
+import ExponentialFamily: ExponentialFamilyDistribution, NormalMeanVariance, getnaturalparameters
+
+function mean(expectation::ClosedWilliamsProduct, f, q::ExponentialFamilyDistribution{T}) where {T <: NormalMeanVariance}
+    η = getnaturalparameters(q)
+    jacobian = @SMatrix [-inv(2*η[2]) η[1]/(2*η[2]^2); 0.0 (-1/η[2])^(3/2)/(2*sqrt(2))]
+    normal = Normal(mean(q), std(q))
+    return mean(expectation, f, normal)' * jacobian
+end

src/exponential_family_interface.jl

@@ -0,0 +1,6 @@
+import ExponentialFamily: ExponentialFamilyDistribution
+
+function mean(expectation::ClosedFormExpectation, f, q::ExponentialFamilyDistribution)
+    dist = convert(Distribution, q)
+    return mean(expectation, f, dist)
+end

test/Exponential/mean_tests.jl

@@ -54,4 +54,20 @@ end
         λ2 = rand(rng)*10
         central_limit_theorem_test(ClosedFormExpectation(), Logpdf(Exponential(λ2)), Exponential(λ1))
     end
+end
+
+@testitem "mean(::ClosedFormExpectation, ::Logpdf{Exponential}, ::ExponentialFamilyDistribution{Exponential}" begin
+    using Distributions
+    using ClosedFormExpectations
+    using ExponentialFamily
+    using StableRNGs
+    using Base.MathConstants: eulergamma
+
+    include("../test_utils.jl")
+    rng = StableRNG(123)
+    for _ in 1:10
+        λ1 = rand(rng)*10
+        λ2 = rand(rng)*10
+        central_limit_theorem_test(ClosedFormExpectation(), Logpdf(Exponential(λ2)), convert(ExponentialFamilyDistribution, Exponential(λ1)))
+    end
 end

test/Normal/williams/ef_parametrization/ef_utils.jl

@@ -0,0 +1,5 @@
+include("../normal_utils.jl")
+
+using ExponentialFamily
+
+score(q::ExponentialFamilyDistribution{NormalMeanVariance}, x) = sufficientstatistics(q, x) .- gradlogpartition(q)

test/Normal/williams/ef_parametrization/williams_tests.jl

@@ -0,0 +1,31 @@
+@testitem "mean(::ClosedWilliamsProduct, p::Logpdf{Normal}, q::ExponentialFamilyDistribution{NormalMeanVariance})" begin
+    include("ef_utils.jl")
+    rng = StableRNG(123)
+    for _ in 1:10
+        μ1, σ1 = rand(rng)*10, rand(rng)*5
+        μ2, σ2 = rand(rng)*10, rand(rng)*5
+        ef = convert(ExponentialFamilyDistribution, Normal(μ1, σ1))
+        central_limit_theorem_test(ClosedWilliamsProduct(), Logpdf(Normal(μ2, σ2)), ef, score)
+    end
+end
+
+@testitem "mean(::ClosedWilliamsProduct, p::Abs, q::ExponentialFamilyDistribution{NormalMeanVariance})" begin
+    include("ef_utils.jl")
+    rng = StableRNG(123)
+    for _ in 1:10
+        μ1, σ1 = rand(rng)*10, rand(rng)*5
+        ef = convert(ExponentialFamilyDistribution, Normal(μ1, σ1))
+        central_limit_theorem_test(ClosedWilliamsProduct(), Abs(), ef, score)
+    end
+end
+
+@testitem "mean(::ClosedWilliamsProduct, p::Logpdf{Laplace}, q::ExponentialFamilyDistribution{NormalMeanVariance})" begin
+    include("ef_utils.jl")
+    rng = StableRNG(123)
+    for _ in 1:10
+        μ1, σ1 = rand(rng)*10, rand(rng)*5
+        loc, θ = rand(rng)*10, rand(rng)*10
+        ef = convert(ExponentialFamilyDistribution, Normal(μ1, σ1))
+        central_limit_theorem_test(ClosedWilliamsProduct(), Logpdf(Laplace(loc, θ)), ef, score)
+    end
+end

test/interface/exponetial_family_interface_tests.jl

@@ -0,0 +1,20 @@
+@testitem "Support ExponetialFamilyDistribution for ClosedFormExpectation" begin
+    include("../test_utils.jl")
+    using ExponentialFamily
+    @testset "mean(::ClosedFormExpectation, f, q::ExponentialFamilyDistribution{Exponential})" begin
+        dist = Exponential(1.0)
+        ef = convert(ExponentialFamilyDistribution, Exponential(1.0))
+        @test mean(ClosedFormExpectation(), Logpdf(Exponential(1.0)), ef) ≈ mean(ClosedFormExpectation(), Logpdf(Exponential(1.0)), dist)
+    end
+    @testset "mean(::ClosedFormExpectation, f, q::ExponentialFamilyDistribution{NormalMeanVariance})" begin
+        dist = Normal(1.0, 1.0)
+        ef = convert(ExponentialFamilyDistribution, Normal(1.0, 1.0))
+        @test mean(ClosedFormExpectation(), Abs(), ef) ≈ mean(ClosedFormExpectation(), Abs(), dist)
+    end
+    @testset "mean(::ClosedFormExpectation, f, q::ExponentialFamilyDistribution{Gamma})" begin
+        import LogExpFunctions: xlogx
+        dist = Gamma(1.0, 1.0)
+        ef = convert(ExponentialFamilyDistribution, Gamma(1.0, 1.0))
+        @test mean(ClosedFormExpectation(), xlogx, ef) ≈ mean(ClosedFormExpectation(), xlogx, dist)
+    end
+end

test/interface/gauss_tests.jl

@@ -27,13 +27,16 @@ end
 
     @testset "ClosedWilliamsProduct interface" begin
         using Distributions
+        using ExponentialFamily
 
         nmv = NormalMeanVariance(0.0, 1.0)
         nmp = NormalMeanPrecision(0.0, 1.0)
+        ef = convert(ExponentialFamilyDistribution, nmv)
 
         @test mean(ClosedWilliamsProduct(), Logpdf(nmv), Normal(0, 1)) isa AbstractArray
         @test mean(ClosedWilliamsProduct(), Logpdf(nmp), Normal(0, 1)) isa AbstractArray
         @test mean(ClosedWilliamsProduct(), Logpdf(nmv), Normal(0, 1)) ≈ mean(ClosedWilliamsProduct(), Logpdf(nmp), Normal(0, 1))
         @test mean(ClosedWilliamsProduct(), Logpdf(nmv), nmv) ≈ mean(ClosedWilliamsProduct(), Logpdf(nmp), nmv)
+        @test mean(ClosedWilliamsProduct(), Logpdf(nmv), ef) isa AbstractArray
     end
 end