fix: change the argument order

Author: Nimrais

src/ClosedFormExpectations.jl

@@ -13,7 +13,7 @@ export ClosedWilliamsProduct
 struct ClosedFormExpectation end
 
 """
-    mean(::ClosedFormExpectation, q, f)
+    mean(::ClosedFormExpectation, f, q)
 
 Compute the E_q[f(x)] where q is a distribution and f is a function.
 """
@@ -22,7 +22,7 @@ function mean(::ClosedFormExpectation, ::Nothing, ::Nothing) end
 struct ClosedWilliamsProduct end 
 
 """
-    mean(::ClosedWilliamsProduct, q, f)
+    mean(::ClosedWilliamsProduct, f, q)
 
 Suppose q is a distribution with density parameterized by θ and f is a function.
 
@@ -44,4 +44,7 @@ include("expressions/Product.jl")
 include("Exponential/Exponential.jl")
 include("Normal/UnivariateNormal.jl")
 
+# Logpdf structure
+include("logpdf.jl")
+
 end

src/Exponential/Exponential.jl

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

src/Normal/UnivariateNormal.jl

@@ -1,13 +1,13 @@
 import Distributions: Laplace, Normal, Rayleigh, params, cdf
 
 
-function mean(::ClosedFormExpectation, q::Normal, p::ComposedFunction{typeof(log), Normal{T}}) where {T}
+function mean(::ClosedFormExpectation, p::ComposedFunction{typeof(log), Normal{T}}, q::Normal) where {T}
     μ_q, σ_q = q.μ, q.σ
     μ_p, σ_p = p.inner.μ, p.inner.σ
     return - 1/2 * log(2 * π * σ_p^2) - (σ_q^2 + (μ_p- μ_q)^2) / (2 * σ_p^2)
 end 
 
-function mean(::ClosedFormExpectation, q::Normal, p::ComposedFunction{typeof(log), Laplace{T}}) where {T}
+function mean(::ClosedFormExpectation, p::ComposedFunction{typeof(log), Laplace{T}}, q::Normal) where {T}
     μ_q, σ_q = q.μ, q.σ
     (μ_p, θ_p) = params(p.inner)
     normal = Normal(0,σ_q)

src/expressions/Product.jl

@@ -4,6 +4,6 @@ struct Product{T} <: Expression
     multipliers::T
 end
 
-function mean(::ClosedFormExpectation, q, p::ComposedFunction{typeof(log), Product{T}}) where {T}
-    return sum(mean(ClosedFormExpectation(), q, log ∘ p_i) for p_i in p.inner.multipliers)
+function mean(::ClosedFormExpectation, p::ComposedFunction{typeof(log), Product{T}}, q) where {T}
+    return sum(mean(ClosedFormExpectation(), log ∘ p_i, q) for p_i in p.inner.multipliers)
 end

src/logpdf.jl

@@ -0,0 +1,15 @@
+export Logpdf
+
+import Distributions: Distribution, logpdf
+
+struct Logpdf{D}
+    dist::D
+end
+
+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

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(), Exponential(λ), log ∘ LogNormal(μ, σ))
+        expectation = mean(ClosedFormExpectation(), log ∘ LogNormal(μ, σ), Exponential(λ))
         @test sigma_rule(expectation, mean(logpdf_samples), std(logpdf_samples), 10^6)
     end
 end
@@ -27,7 +27,7 @@ end
         N = 10^6
         samples = rand(rng, Exponential(λ), 10^6)
         log_samples = log.(samples)
-        @test sigma_rule(mean(ClosedFormExpectation(), Exponential(λ), log), mean(log_samples), std(log_samples), N)
+        @test sigma_rule(mean(ClosedFormExpectation(), log, Exponential(λ)), mean(log_samples), std(log_samples), N)
     end
 end
 
@@ -46,7 +46,7 @@ end
         N = 10^5
         samples = rand(rng, Exponential(λ), N)
         log_samples = log.(ExpLogSquare(μ, σ).(samples))
-        @test sigma_rule(mean(ClosedFormExpectation(), Exponential(λ), log ∘ ExpLogSquare(μ, σ)), mean(log_samples), std(log_samples), N)
+        @test sigma_rule(mean(ClosedFormExpectation(), log ∘ ExpLogSquare(μ, σ), Exponential(λ)), mean(log_samples), std(log_samples), N)
     end
 end
 
@@ -63,8 +63,8 @@ end
         σ = rand(rng)*10
         λ = rand(rng)*10
         product = ClosedFormExpectations.Product((ExpLogSquare(μ, σ), identity))
-        sum_mean = mean(ClosedFormExpectation(), Exponential(λ), log ∘ ExpLogSquare(μ, σ)) + mean(ClosedFormExpectation(), Exponential(λ), log)
-        @test mean(ClosedFormExpectation(), Exponential(λ), log ∘ product) ≈ sum_mean
+        sum_mean = mean(ClosedFormExpectation(), log ∘ ExpLogSquare(μ, σ), Exponential(λ)) + mean(ClosedFormExpectation(), log, Exponential(λ))
+        @test mean(ClosedFormExpectation(), log ∘ product, Exponential(λ)) ≈ sum_mean
     end
 end
 
@@ -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(), Exponential(λ1), log ∘ Exponential(λ2)), mean(log_samples), std(log_samples), N)
+        @test sigma_rule(mean(ClosedFormExpectation(), log ∘ Exponential(λ2), Exponential(λ1)), mean(log_samples), std(log_samples), N)
     end
 end

test/Exponential/williams_tests.jl

@@ -16,7 +16,7 @@
         N = 10^6
         samples = rand(rng, Exponential(λ), 10^6)
         williams_product = map(x -> score(Exponential(λ), x)*log(x), samples)
-        expectation = mean(ClosedWilliamsProduct(), Exponential(λ), identity)
+        expectation = mean(ClosedWilliamsProduct(), identity, Exponential(λ))
         @test sigma_rule(expectation, mean(williams_product), std(williams_product), N)
     end
 end
@@ -42,7 +42,7 @@ end
         samples = rand(rng, Exponential(λ), 10^6)
         fn(x)  = (log ∘ ExpLogSquare(μ, σ))(x)
         williams_product = map(x -> score(Exponential(λ), x)*fn(x), samples)
-        expectation = mean(ClosedWilliamsProduct(), Exponential(λ), ExpLogSquare(μ, σ))
+        expectation = mean(ClosedWilliamsProduct(), ExpLogSquare(μ, σ), Exponential(λ))
         @test sigma_rule(expectation, mean(williams_product), std(williams_product), N)
     end
 end
@@ -68,7 +68,7 @@ end
         samples = rand(rng, Exponential(λ), 10^6)
         fn(x)  = logpdf(LogNormal(μ, σ), x)
         williams_product = map(x -> score(Exponential(λ), x)*fn(x), samples)
-        expectation = mean(ClosedWilliamsProduct(), Exponential(λ), LogNormal(μ, σ))
+        expectation = mean(ClosedWilliamsProduct(), LogNormal(μ, σ), Exponential(λ))
         @test sigma_rule(expectation, mean(williams_product), std(williams_product), N)
     end
 end
@@ -90,6 +90,6 @@ end
         samples = rand(rng, Exponential(λ1), N)
         fn(x) = logpdf(Exponential(λ2), x)
         williams_product = map(x -> score(Exponential(λ1), x)*fn(x), samples)
-        @test sigma_rule(mean(ClosedWilliamsProduct(), Exponential(λ1), Exponential(λ2)), mean(williams_product), std(williams_product), N)
+        @test sigma_rule(mean(ClosedWilliamsProduct(), Exponential(λ2), Exponential(λ1)), mean(williams_product), std(williams_product), N)
     end
 end

test/Normal/mean_test.jl

@@ -11,7 +11,7 @@
         N = 10^5
         samples = rand(rng, Normal(μ_1, σ_1), N)
         log_samples = map(x -> logpdf(Normal(μ_2, σ_2),x),samples)
-        @test sigma_rule(mean(ClosedFormExpectation(), Normal(μ_1, σ_1), log ∘ Normal(μ_2, σ_2)), mean(log_samples), std(log_samples), N)
+        @test sigma_rule(mean(ClosedFormExpectation(), log ∘ Normal(μ_2, σ_2), Normal(μ_1, σ_1)), mean(log_samples), std(log_samples), N)
     end
 end
 
@@ -30,6 +30,6 @@ end
         laplace = Laplace(μ_2, θ)
         samples = rand(rng, normal, N)
         log_samples = map(x -> logpdf(laplace,x),samples)
-        @test sigma_rule(mean(ClosedFormExpectation(), normal, log ∘ laplace), mean(log_samples), std(log_samples), N)
+        @test sigma_rule(mean(ClosedFormExpectation(), log ∘ laplace, normal), mean(log_samples), std(log_samples), N)
     end
 end

test/interface/fix1.jl

@@ -0,0 +1,18 @@
+@testset "Support ::Base.Fix1{typeof(logpdf), D}" begin
+    using Distributions
+    using ClosedFormExpectations
+    using StableRNGs
+    using Base.MathConstants: eulergamma
+
+    include("../test_utils.jl")
+    rng = StableRNG(123)
+    for _ in 1:10
+        μ = rand(rng)*10
+        σ = rand(rng)*10
+        λ = rand(rng)*10
+        N = 10^5
+        samples = rand(rng, Exponential(λ), N)
+        log_samples = log.(ExpLogSquare(μ, σ).(samples))
+        @test sigma_rule(mean(ClosedFormExpectation(), Exponential(λ), log ∘ ExpLogSquare(μ, σ)), mean(log_samples), std(log_samples), N)
+    end
+end