Specialize logdensityof for most combinators

Improves type stability and numerical type propagation.
Specialize logdensityof until a full refactor of the density engine.

Author: oschulz

src/MeasureBase.jl

@@ -109,6 +109,7 @@ function logdensity_def end
 using Compat
 
 using IrrationalConstants
+using IrrationalConstants: loghalf
 
 include("static.jl")
 include("smf.jl")

src/combinators/half.jl

@@ -19,6 +19,11 @@ function Base.rand(rng::AbstractRNG, ::Type{T}, μ::Half) where {T}
     return abs(rand(rng, T, unhalf(μ)))
 end
 
+function logdensityof(μ::Half, x)
+    ld = logdensityof(unhalf(μ), x) - loghalf
+    return x ≥ 0 ? ld : oftype(ld, -Inf)
+end
+
 logdensity_def(μ::Half, x) = logdensity_def(unhalf(μ), x)
 
 @inline function insupport(d::Half, x)

src/combinators/power.jl

@@ -78,25 +78,27 @@ params(d::PowerMeasure) = params(first(marginals(d)))
     basemeasure(d.parent)^d.axes
 end
 
-@inline function logdensity_def(d::PowerMeasure{M}, x) where {M}
-    parent = d.parent
-    sum(x) do xj
-        logdensity_def(parent, xj)
+for func in [:logdensityof, :logdensity_def]
+    @eval @inline function $func(d::PowerMeasure{M}, x) where {M}
+        parent = d.parent
+        sum(x) do xj
+            $func(parent, xj)
+        end
     end
-end
 
-@inline function logdensity_def(d::PowerMeasure{M,Tuple{Static.SOneTo{N}}}, x) where {M,N}
-    parent = d.parent
-    sum(1:N) do j
-        @inbounds logdensity_def(parent, x[j])
+    @eval @inline function $func(d::PowerMeasure{M,Tuple{Static.SOneTo{N}}}, x) where {M,N}
+        parent = d.parent
+        sum(1:N) do j
+            @inbounds $func(parent, x[j])
+        end
     end
-end
 
-@inline function logdensity_def(
-    d::PowerMeasure{M,NTuple{N,Static.SOneTo{0}}},
-    x,
-) where {M,N}
-    static(0.0)
+    @eval @inline function $func(
+        d::PowerMeasure{M,NTuple{N,Static.SOneTo{0}}},
+        x,
+    ) where {M,N}
+        static(0.0)
+    end
 end
 
 @inline function insupport(μ::PowerMeasure, x)

src/combinators/product.jl

@@ -72,8 +72,10 @@ function _rand_product(
     end |> collect
 end
 
-@inline function logdensity_def(d::AbstractProductMeasure, x)
-    mapreduce(logdensity_def, +, marginals(d), x)
+for func in [:logdensityof, :logdensity_def]
+    @eval @inline function $func(d::AbstractProductMeasure, x)
+        mapreduce($func, +, marginals(d), x)
+    end
 end
 
 struct ProductMeasure{M} <: AbstractProductMeasure
@@ -88,27 +90,37 @@ function Pretty.tile(d::ProductMeasure{T}) where {T<:Tuple}
     Pretty.list_layout(Pretty.tile.([marginals(d)...]), sep = " ⊗ ")
 end
 
-# For tuples, `mapreduce` has trouble with type inference
-@inline function logdensity_def(d::ProductMeasure{T}, x) where {T<:Tuple}
-    ℓs = map(logdensity_def, marginals(d), x)
-    sum(ℓs)
-end
-
-@generated function logdensity_def(d::ProductMeasure{NamedTuple{N,T}}, x) where {N,T}
+@eval @generated function _product_gen_impl(
+    ::Val{func},
+    d::ProductMeasure{NamedTuple{N,T}},
+    x,
+) where {func,N,T}
     k1 = QuoteNode(first(N))
     q = quote
         m = marginals(d)
-        ℓ = logdensity_def(getproperty(m, $k1), getproperty(x, $k1))
+        ℓ = $func(getproperty(m, $k1), getproperty(x, $k1))
     end
     for k in Base.tail(N)
         k = QuoteNode(k)
-        qk = :(ℓ += logdensity_def(getproperty(m, $k), getproperty(x, $k)))
+        qk = :(ℓ += $func(getproperty(m, $k), getproperty(x, $k)))
         push!(q.args, qk)
     end
 
     return q
 end
 
+for func in [:logdensityof, :logdensity_def]
+    # For tuples, `mapreduce` has trouble with type inference
+    @eval @inline function $func(d::ProductMeasure{T}, x) where {T<:Tuple}
+        ℓs = map($func, marginals(d), x)
+        sum(ℓs)
+    end
+
+    @eval function $func(d::ProductMeasure{NamedTuple{N,T}}, x) where {N,T}
+        _product_gen_impl(Val($func), d, x)
+    end
+end
+
 # @generated function basemeasure(d::ProductMeasure{NamedTuple{N,T}}, x) where {N,T}
 #     q = quote
 #         m = marginals(d)

src/combinators/spikemixture.jl

@@ -21,11 +21,17 @@ end
     SpikeMixture(basemeasure(μ.m), static(1.0), static(1.0))
 end
 
-@inline function logdensity_def(μ::SpikeMixture, x)
-    if iszero(x)
-        return log(μ.s)
-    else
-        return log(μ.w) + logdensity_def(μ.m, x)
+for func in [:logdensityof, :logdensity_def]
+    @eval @inline function $func(μ::SpikeMixture, x)
+        R1 = Core.Compiler.return_type(log, Tuple{typeof(μ.s)})
+        R2 = Core.Compiler.return_type(log, Tuple{typeof(μ.w)})
+        R3 = Core.Compiler.return_type($func, Tuple{typeof(μ.m),typeof(x)})
+        R = promote_type(R1, R2, R3)
+        if iszero(x)
+            return convert(R, log(μ.s))::R
+        else
+            return convert(R, log(μ.w) + $func(μ.m, x))::R
+        end
     end
 end
 

src/combinators/transformedmeasure.jl

@@ -55,27 +55,29 @@ end
 #     logdensity_rel(pushfwd(f, inv_f, ν.origin, WithVolCorr()), β.origin, x)
 # end
 
-@inline function logdensity_def(ν::PushforwardMeasure{F,I,M,<:WithVolCorr}, y) where {F,I,M}
-    f = ν.f
-    finv = ν.finv
-    x_orig, inv_ladj = with_logabsdet_jacobian(unwrap(finv), y)
-    logd_orig = logdensity_def(ν.origin, x_orig)
-    logd = float(logd_orig + inv_ladj)
-    neginf = oftype(logd, -Inf)
-    return ifelse(
-        # Zero density wins against infinite volume:
-        (isnan(logd) && logd_orig == -Inf && inv_ladj == +Inf) ||
-        # Maybe  also for (logd_orig == -Inf) && isfinite(inv_ladj) ?
-        # Return constant -Inf to prevent problems with ForwardDiff:
-        (isfinite(logd_orig) && (inv_ladj == -Inf)),
-        neginf,
-        logd,
-    )
-end
+for func in [:logdensityof, :logdensity_def]
+    @eval @inline function $func(ν::PushforwardMeasure{F,I,M,<:WithVolCorr}, y) where {F,I,M}
+        f = ν.f
+        finv = ν.finv
+        x_orig, inv_ladj = with_logabsdet_jacobian(unwrap(finv), y)
+        logd_orig = $func(ν.origin, x_orig)
+        logd = float(logd_orig + inv_ladj)
+        neginf = oftype(logd, -Inf)
+        return ifelse(
+            # Zero density wins against infinite volume:
+            (isnan(logd) && logd_orig == -Inf && inv_ladj == +Inf) ||
+            # Maybe  also for (logd_orig == -Inf) && isfinite(inv_ladj) ?
+            # Return constant -Inf to prevent problems with ForwardDiff:
+            (isfinite(logd_orig) && (inv_ladj == -Inf)),
+            neginf,
+            logd,
+        )
+    end
 
-@inline function logdensity_def(ν::PushforwardMeasure{F,I,M,<:NoVolCorr}, y) where {F,I,M}
-    x = ν.finv(y)
-    return logdensity_def(ν.origin, x)
+    @eval @inline function $func(ν::PushforwardMeasure{F,I,M,<:NoVolCorr}, y) where {F,I,M}
+        x = ν.finv(y)
+        return $func(ν.origin, x)
+    end
 end
 
 insupport(ν::PushforwardMeasure, y) = insupport(ν.origin, ν.finv(y))