refactor: LogExpression -> Expression

Author: Nimrais

src/ClosedFormExpectations.jl

@@ -30,14 +30,15 @@ Compute the E_q[f(x) ∇_θ log q(x; θ)] where q is a distribution and f is a f
 """
 function mean(::ClosedWilliamsProduct, ::Nothing, ::Nothing) end
 
-abstract type LogExpression end
+abstract type Expression end
 
-function (f::ComposedFunction{typeof(log), T})(x) where {T <: LogExpression}
+function (f::ComposedFunction{typeof(log), T})(x) where {T <: Expression}
     return log(f.inner, x)
 end
 
 # expressions
 include("expressions/ExpLogSquare.jl")
+include("expressions/Product.jl")
 
 # rules for computing expectation of log
 include("Exponential/Exponential.jl")

src/Exponential/Exponential.jl

@@ -11,6 +11,10 @@ function mean(::ClosedFormExpectation, q::Exponential, p::typeof(log))
     return -eulergamma + log(mean(q))
 end
 
+function mean(::ClosedFormExpectation, q::Exponential, p::ComposedFunction{typeof(log), typeof(identity)})
+    return -eulergamma + log(mean(q))
+end
+
 function mean(::ClosedFormExpectation, q::Exponential, p::ComposedFunction{typeof(log), ExpLogSquare{T}}) where {T}
     μ, σ = p.inner.μ, p.inner.σ
     λ = mean(q)

src/expressions/ExpLogSquare.jl

@@ -3,7 +3,7 @@
 
 ExpLogSquare is a type that represents the exp(-(log(x) - μ)^2/(2σ^2)) function.
 """
-struct ExpLogSquare{T} <: LogExpression
+struct ExpLogSquare{T} <: Expression
     μ::T
     σ::T
 end

src/expressions/Product.jl

@@ -1,7 +1,9 @@
-struct Product <: LogExpression
-    multipliers::Tuple{<:LogExpression}
+export Product
+
+struct Product{T} <: Expression
+    multipliers::T
 end
 
-function mean(::ClosedFormExpectation, q, p::ComposedFunction{typeof(log), Product})
-    return sum(mean(ClosedFormExpectation(), q, log ∘ p_i) for p_i in p.multipliers)
+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)
 end

test/Exponential/mean_tests.jl

@@ -31,7 +31,7 @@ end
     end
 end
 
-@testitem "mean(::Exponential, ::ComposedFunction{typeof(log), ExpLogSquare}" begin
+@testitem "mean(::Exponential, ::ComposedFunction{typeof(log), ExpLogSquare})" begin
     using Distributions
     using ClosedFormExpectations
     using StableRNGs
@@ -48,4 +48,22 @@ end
         log_samples = log.(ExpLogSquare(μ, σ).(samples))
         @test sigma_rule(mean(ClosedFormExpectation(), Exponential(λ), log ∘ ExpLogSquare(μ, σ)), mean(log_samples), std(log_samples), N)
     end
+end
+
+@testitem "mean(::Exponential, ::ComposedFunction{typeof(log), ExpLogSquare x identity}" 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
+        product = ClosedFormExpectations.Product((ExpLogSquare(μ, σ), identity))
+        sum_mean = mean(ClosedFormExpectation(), Exponential(λ), log ∘ ExpLogSquare(μ, σ)) + mean(ClosedFormExpectation(), Exponential(λ), log)
+        @test mean(ClosedFormExpectation(), Exponential(λ), log ∘ product) ≈ sum_mean
+    end
 end