Specialize logdensityof for primitive measures

We need maximum performance on these.

Also try to preserve floating point types as far as possible.

Author: oschulz

src/primitive.jl

@@ -19,6 +19,9 @@ basemeasure(μ::PrimitiveMeasure) = μ
 
 @inline basemeasure_depth(::PrimitiveMeasure) = static(0)
 
+@inline logdensityof(::PrimitiveMeasure, x::Real) = zero(float(typeof(x)))
+@inline logdensityof(::PrimitiveMeasure, x) = false
+
 logdensity_def(::PrimitiveMeasure, x) = static(0.0)
 
 logdensity_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0

src/primitives/counting.jl

@@ -12,10 +12,15 @@ struct Counting{T} <: AbstractMeasure
     Counting(supp) = new{Core.Typeof(supp)}(supp)
 end
 
-function logdensity_def(μ::Counting, x)
-    insupport(μ, x) ? 0.0 : -Inf
+@inline function logdensityof(μ::Counting, x::Real)
+    R = float(typeof(x))
+    insupport(μ, x) ? zero(R) : R(-Inf)
 end
 
+@inline logdensityof(μ::Counting, x) = insupport(μ, x) ? 0.0 : -Inf
+
+@inline logdensity_def(μ::Counting, x) = logdensityof(μ, x)
+
 basemeasure(::Counting) = CountingBase()
 
 Counting() = Counting(ℤ)

src/primitives/dirac.jl

@@ -20,7 +20,15 @@ basemeasure(d::Dirac) = CountingBase()
 
 massof(::Dirac) = static(1.0)
 
-logdensity_def(μ::Dirac, x) = 0.0
+function logdensityof(μ::Dirac, x::Real)
+    R = float(typeof(x))
+    insupport(μ, x) ? zero(R) : R(-Inf)
+end
+
+logdensityof(μ::Dirac, x) = insupport(μ, x) ? 0.0 : -Inf
+
+logdensity_def(::Dirac, x::Real) = zero(float(typeof(x)))
+logdensity_def(::Dirac, x) = 0.0
 
 Base.rand(::Random.AbstractRNG, T::Type, μ::Dirac) = μ.x
 

src/primitives/lebesgue.jl

@@ -63,6 +63,13 @@ insupport(μ::Lebesgue, x) = x ∈ μ.support
 
 insupport(::Lebesgue{RealNumbers}, ::Real) = true
 
+@inline function logdensityof(μ::Lebesgue, x::Real)
+    R = float(typeof(x))
+    insupport(μ, x) ? zero(R) : R(-Inf)
+end
+
+@inline logdensityof(μ::Lebesgue, x) = insupport(μ, x) ? 0.0 : -Inf
+
 massof(::Lebesgue{RealNumbers}, s::Interval) = width(s)
 
 # Example: 

test/runtests.jl

@@ -11,6 +11,7 @@ include("test_aqua.jl")
 
 include("static.jl")
 
+include("test_primitive.jl")
 include("test_standard.jl")
 include("test_basics.jl")
 

test/test_primitive.jl

@@ -0,0 +1,39 @@
+using Test
+
+using MeasureBase
+using MeasureBase: insupport as measure_insupport
+
+using DensityInterface: logdensityof
+
+@testset "primitive" begin
+    for (m, x) in [
+        (MeasureBase.LebesgueBase(), -1.0f0),
+        (MeasureBase.LebesgueBase(), 1.0f0),
+        (Lebesgue(), 1.0f0),
+        (Lebesgue(), -1.0f0),
+        (MeasureBase.CountingBase(), -1.0f0),
+        (MeasureBase.CountingBase(), 1.0f0),
+        (Counting(), 2),
+        (Counting(), 2.0f0),
+        (Counting(), 1.5f0),
+        (Dirac(4.2), 4.2f0),
+        (Dirac(4.2), -1.0f0),
+        (Dirac([1, 2, 3]), [1, 2, 3]),
+        (Dirac([4, 5]), [4, 5]),
+    ]
+        @testset "$(nameof(typeof(m)))" begin
+            for x in [-Inf, -1.2, -1, 0, 0.0, 1 // 2, 0.5, 1, 1.0, 2, 2.3]
+                if measure_insupport(m, x)
+                    @test @inferred(logdensityof(m, x)) ≈ 0
+                    if x isa Real
+                        ld = logdensityof(m, x)
+                        ld isa float(typeof(x))
+                    end
+                    @test MeasureBase.unsafe_logdensityof(m, x) ≈ 0
+                else
+                    @test @inferred(logdensityof(m, x)) ≈ -Inf
+                end
+            end
+        end
+    end
+end