fix: support Base.Fix1(logpdf, dist)

Author: Nimrais

README.md

@@ -23,5 +23,5 @@ The package exports the following:
 ```julia
 using ClosedFormExpectations
 using Distributions
-mean(ClosedFormExpectation(), Exponetial(10), log)
+mean(ClosedFormExpectation(), log, Exponetial(10))
 ```

src/ClosedFormExpectations.jl

@@ -36,6 +36,9 @@ function (f::ComposedFunction{typeof(log), T})(x) where {T <: Expression}
     return log(f.inner, x)
 end
 
+# Logpdf structure
+include("logpdf.jl")
+
 # expressions
 include("expressions/ExpLogSquare.jl")
 include("expressions/Product.jl")
@@ -44,7 +47,6 @@ include("expressions/Product.jl")
 include("Exponential/Exponential.jl")
 include("Normal/UnivariateNormal.jl")
 
-# Logpdf structure
-include("logpdf.jl")
+
 
 end

src/Exponential/Exponential.jl

@@ -2,14 +2,14 @@ import Distributions: Exponential, LogNormal, scale, kldivergence, entropy
 import Base.MathConstants: eulergamma
 import LogExpFunctions: xlogx
 
-function mean(::ClosedFormExpectation, p::ComposedFunction{typeof(log), Exponential{T}}, q::Exponential) where {T}
+function mean(::ClosedFormExpectation, p::Logpdf{Exponential{T}}, q::Exponential) where {T}
     λ1 = mean(q)
-    λ2 = mean(p.inner)
+    λ2 = mean(p.dist)
     return -(λ1 + xlogx(λ2))/λ2
 end
 
-function mean(::ClosedFormExpectation, p::ComposedFunction{typeof(log), LogNormal{T}}, q::Exponential) where {T}
-    μ, σ = p.inner.μ, p.inner.σ
+function mean(::ClosedFormExpectation, p::Logpdf{LogNormal{T}}, q::Exponential) where {T}
+    μ, σ = p.dist.μ, p.dist.σ
     λ = mean(q)
     return 1/(2*σ^2)*(-(μ+eulergamma)^2 - π^2/6 - log(λ)*(-2*(eulergamma+μ) + log(λ))) + eulergamma - log(λ) - 0.5*log(2π) - log(σ)
 end 

src/logpdf.jl

@@ -10,6 +10,4 @@ function (f::Logpdf{D})(args...) where {D <: Distribution}
     return logpdf(f.dist, args...;)
 end
 
-# function convert(::Type{Logpdf}, fixed_call::Fix1{typeof{logpdf}, D}) where {D <: Distribution}
-#     return Logpdf(D)
-# end
+mean(::ClosedFormExpectation, f::Base.Fix1{typeof(logpdf), D}, q) where {D} = mean(ClosedFormExpectation(), Logpdf(f.x), q)

test/Exponential/mean_tests.jl

@@ -11,7 +11,7 @@
         N = 10^6
         samples = rand(rng, Exponential(λ), N)
         logpdf_samples = logpdf(LogNormal(μ, σ), samples)        
-        expectation = mean(ClosedFormExpectation(), log ∘ LogNormal(μ, σ), Exponential(λ))
+        expectation = mean(ClosedFormExpectation(), Logpdf(LogNormal(μ, σ)), Exponential(λ))
         @test sigma_rule(expectation, mean(logpdf_samples), std(logpdf_samples), 10^6)
     end
 end
@@ -68,7 +68,7 @@ end
     end
 end
 
-@testitem "mean(::Exponential, ::ComposedFunction{typeof(log), Exponential}" begin
+@testitem "mean(::ClosedFormExpectation, ::Logpdf{Exponential}, ::Exponential)" begin
     using Distributions
     using ClosedFormExpectations
     using StableRNGs
@@ -82,6 +82,6 @@ end
         N = 10^6
         samples = rand(rng, Exponential(λ1), N)
         log_samples = logpdf(Exponential(λ2), samples)
-        @test sigma_rule(mean(ClosedFormExpectation(), log ∘ Exponential(λ2), Exponential(λ1)), mean(log_samples), std(log_samples), N)
+        @test sigma_rule(mean(ClosedFormExpectation(), Logpdf(Exponential(λ2)), Exponential(λ1)), mean(log_samples), std(log_samples), N)
     end
 end