Specialize and test logdensityof for standard measures

We need maximum performance for these.

Author: oschulz

src/standard/stdexponential.jl

@@ -2,7 +2,12 @@ struct StdExponential <: StdMeasure end
 
 export StdExponential
 
-insupport(d::StdExponential, x) = x ≥ zero(x)
+insupport(::StdExponential, x) = x ≥ zero(x)
+
+@inline function logdensityof(::StdExponential, x)
+    R = float(typeof(x))
+    x ≥ zero(R) ? convert(R, -x) : R(-Inf)
+end
 
 @inline logdensity_def(::StdExponential, x) = -x
 @inline basemeasure(::StdExponential) = LebesgueBase()

src/standard/stdlogistic.jl

@@ -4,7 +4,9 @@ export StdLogistic
 
 @inline insupport(d::StdLogistic, x) = true
 
-@inline logdensity_def(::StdLogistic, x) = (u = -abs(x); u - 2 * log1pexp(u))
+@inline logdensityof(::StdLogistic, x) = (u = -abs(x); u - 2 * log1pexp(u))
+
+@inline logdensity_def(::StdLogistic, x) = logdensityof(StdLogistic(), x)
 @inline basemeasure(::StdLogistic) = LebesgueBase()
 
 @inline transport_def(::StdUniform, μ::StdLogistic, x) = logistic(x)

src/standard/stdnormal.jl

@@ -1,11 +1,13 @@
 using SpecialFunctions: erfc, erfcinv
-using IrrationalConstants: invsqrt2
+using IrrationalConstants: invsqrt2, log2π
 
 struct StdNormal <: StdMeasure end
 
 export StdNormal
 
-@inline insupport(d::StdNormal, x) = true
+@inline insupport(::StdNormal, x) = true
+
+@inline logdensityof(::StdNormal, x) = (-x^2 - log2π) / 2
 
 @inline logdensity_def(::StdNormal, x) = -x^2 / 2
 @inline basemeasure(::StdNormal) = WeightedMeasure(static(-0.5 * log2π), LebesgueBase())

src/standard/stduniform.jl

@@ -2,7 +2,12 @@ struct StdUniform <: StdMeasure end
 
 export StdUniform
 
-insupport(d::StdUniform, x) = zero(x) ≤ x ≤ one(x)
+insupport(::StdUniform, x) = zero(x) ≤ x ≤ one(x)
+
+@inline function logdensityof(::StdUniform, x)
+    R = float(typeof(x))
+    zero(x) ≤ x ≤ one(x) ? zero(R) : R(-Inf)
+end
 
 @inline logdensity_def(::StdUniform, x) = zero(x)
 @inline basemeasure(::StdUniform) = LebesgueBase()

test/Project.toml

@@ -3,6 +3,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"

test/runtests.jl

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

test/test_basics.jl

@@ -124,13 +124,13 @@ end
     @test logdensityof(Lebesgue()^3, 2) == logdensityof(Lebesgue()^(3, 1), (2, 0))
 end
 
-Normal() = ∫exp(x -> -0.5x^2, Lebesgue(ℝ))
+NormalMeasure() = ∫exp(x -> -0.5x^2, Lebesgue(ℝ))
 
 @testset "Half" begin
-    HalfNormal() = Half(Normal())
+    HalfNormal() = Half(NormalMeasure())
     @test logdensityof(HalfNormal(), -0.2) == -Inf
-    @test logdensity_def(HalfNormal(), 0.2) == logdensity_def(Normal(), 0.2)
-    @test densityof(HalfNormal(), 0.2) ≈ 2 * densityof(Normal(), 0.2)
+    @test logdensity_def(HalfNormal(), 0.2) == logdensity_def(NormalMeasure(), 0.2)
+    @test densityof(HalfNormal(), 0.2) ≈ 2 * densityof(NormalMeasure(), 0.2)
 end
 
 @testset "Likelihood" begin
@@ -218,7 +218,7 @@ end
         @test log(f(x)) ≈ x^2
     end
 
-    let f = log𝒹(∫exp(x -> x^2, Normal()), Normal())
+    let f = log𝒹(∫exp(x -> x^2, NormalMeasure()), NormalMeasure())
         @test f(x) ≈ x^2
     end
 end

test/test_standard.jl

@@ -0,0 +1,28 @@
+using Test
+using MeasureBase
+using MeasureBase: insupport as measure_insupport
+
+using DensityInterface: logdensityof
+import Distributions: insupport as dist_insupport
+using Distributions: Normal, Exponential, Logistic, Uniform
+
+@testset "standard" begin
+    for (m, d) in [
+        (StdUniform(), Uniform()),
+        (StdExponential(), Exponential()),
+        (StdLogistic(), Logistic()),
+        (StdNormal(), Normal()),
+    ]
+        @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]
+                @test @inferred(logdensityof(m, x)) ≈ logdensityof(d, x)
+                @test @inferred(logdensityof(m, x)) ≈ logdensity_rel(m, rootmeasure(m), x)
+                @test measure_insupport(m, x) ≈ dist_insupport(d, x)
+                if measure_insupport(m, x)
+                    @test @inferred(MeasureBase.unsafe_logdensityof(m, x)) ≈
+                          logdensityof(d, x)
+                end
+            end
+        end
+    end
+end