Dev (#15)

* update 3-arg logdensity

* update logjac

* factoredbase

* update power measure

* update half.jl to use Factoredbase

* drop outdated test

* add a test

* bump version

Author: cscherrer

Project.toml

@@ -1,7 +1,7 @@
 name = "MeasureBase"
 uuid = "fa1605e6-acd5-459c-a1e6-7e635759db14"
 authors = ["Chad Scherrer <[email protected]> and contributors"]
-version = "0.3.7"
+version = "0.4.0"
 
 [deps]
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"

src/MeasureBase.jl

@@ -49,6 +49,7 @@ include("primitives/lebesgue.jl")
 include("primitives/dirac.jl")
 include("primitives/trivial.jl")
 
+include("combinators/factoredbase.jl")
 include("combinators/weighted.jl")
 include("combinators/affine.jl")
 include("combinators/superpose.jl")

src/combinators/affine.jl

@@ -22,6 +22,13 @@ Base.propertynames(d::AffineTransform{N}) where {N} = N
 (f::AffineTransform{(:μ,:σ)})(x) = f.σ * x + f.μ
 (f::AffineTransform{(:μ,:ω)})(x) = f.ω \ x + f.μ
 
+# TODO: `log` doesn't work for the multivariate case, we need the log absolute determinant
+logjac(f::AffineTransform{(:μ,:σ)}) = log(f.σ)
+logjac(f::AffineTransform{(:μ,:ω)}) = -log(f.ω)
+logjac(f::AffineTransform{(:σ,)}) = log(f.σ)
+logjac(f::AffineTransform{(:ω,)}) = -log(f.ω)
+logjac(f::AffineTransform{(:μ,)}) = 0.0
+
 ###############################################################################
 
 struct Affine{N,M,T} <: AbstractMeasure
@@ -61,8 +68,6 @@ Base.propertynames(d::Affine{N}) where {N} = N ∪ (:parent,)
     end
 end
 
-
-
 # Note: We could also write
 # logdensity(d::Affine, x) = logdensity(inv(getfield(d, :f)), x)
 
@@ -74,16 +79,12 @@ logdensity(d::Affine{(:μ,)}, x) = logdensity(d.parent, x - d.μ)
 
 basemeasure(d::Affine) = Affine(getfield(d, :f), basemeasure(d.parent))
 
-basemeasure(d::Affine{N,L}) where {N, L<:Lebesgue} = d.parent
+# We can't do this until we know we're working with Lebesgue measure, since for
+# example it wouldn't make sense to apply a log-Jacobian to a point measure
+basemeasure(d::Affine{N,L}) where {N, L<:Lebesgue} = WeightedMeasure(-logjac(d), d.parent)
 
-logdensity(d::Affine{N,L}, x) where {N,L<:Lebesgue} = logjac(getfield(d, :f))
+logjac(d::Affine) = logjac(getfield(d, :f))
 
-# TODO: `log` doesn't work for the multivariate case, we need the log absolute determinant
-logjac(::AffineTransform{(:μ,:σ)}) = -log(d.σ)
-logjac(::AffineTransform{(:μ,:ω)}) = log(d.ω)
-logjac(::AffineTransform{(:σ,)}) = -log(d.σ)
-logjac(::AffineTransform{(:ω,)}) = log(d.ω)
-logjac(::AffineTransform{(:μ,)}) = 0.0
 
 function Base.rand(rng::Random.AbstractRNG, ::Type{T}, d::Affine) where {T}
     z = rand(rng, T, parent(d))

src/combinators/factoredbase.jl

@@ -0,0 +1,15 @@
+export FactoredBase
+
+struct FactoredBase{R,C,V,B} <: AbstractMeasure
+    inbounds :: R
+    constℓ :: C
+    varℓ :: V
+    base :: B
+end
+
+function logdensity(d::FactoredBase, x)
+    d.inbounds(x) || return -Inf
+    d.constℓ + d.varℓ
+end
+
+basemeasure(d::FactoredBase) = d.base

src/combinators/half.jl

@@ -11,10 +11,16 @@ end
 
 unhalf(μ::Half) = μ.parent
 
-basemeasure(μ::Half) = WeightedMeasure(logtwo, basemeasure(unhalf(μ)))
+function basemeasure(μ::Half)
+    inbounds(x) = x > 0
+    constℓ = logtwo
+    varℓ = 0.0
+    base = basemeasure(unhalf(μ))
+    FactoredBase(inbounds, constℓ, varℓ, base)
+end
 
 function Base.rand(rng::AbstractRNG, T::Type, μ::Half)
     return abs(rand(rng, T, unhalf(μ)))
 end
 
-logdensity(μ::Half, x) = x > 0 ? logdensity(unhalf(μ), x) : -Inf
+logdensity(μ::Half, x) = logdensity(unhalf(μ), x)

src/combinators/power.jl

@@ -63,3 +63,21 @@ params(d::ProductMeasure{F,<:Fill}) where {F} = params(first(marginals(d)))
 params(::Type{P}) where {F,P<:ProductMeasure{F,<:Fill}} = params(D)
 
 # basemeasure(μ::PowerMeasure) = @inbounds basemeasure(first(μ.data))^size(μ.data)
+
+@inline basemeasure(d::PowerMeasure) = _basemeasure(d, (basemeasure(d.f(first(d.pars)))))
+
+@inline _basemeasure(d::PowerMeasure, b) = b ^ size(d.pars)
+
+@inline function _basemeasure(d::PowerMeasure, b::FactoredBase)
+    n = length(d.pars)
+    inbounds(x) = all(xj -> b.inbounds(xj), x)
+    constℓ = n * b.constℓ
+    varℓ = n * b.varℓ
+    base = b.base ^ size(d.pars)
+    FactoredBase(inbounds, constℓ, varℓ, base)
+end
+
+function logdensity(d::PowerMeasure, x)
+    d1 = d.f(first(d.pars))
+    sum(xj -> logdensity(d1, xj), x)
+end

src/density.jl

@@ -107,6 +107,8 @@ function logdensity(μ::AbstractMeasure, ν::AbstractMeasure, x)
     α = basemeasure(μ)
     β = basemeasure(ν)
 
+    # If α===μ and β===ν, The recursive call would be exactly the same as the
+    # original one. We need to break the recursion.
     if α===μ && β===ν
         @warn """
         No method found for logdensity(μ, ν, x) where
@@ -119,10 +121,18 @@ function logdensity(μ::AbstractMeasure, ν::AbstractMeasure, x)
         return NaN
     end
 
-    ℓμ = logdensity(μ, x)
-    ℓν = logdensity(ν, x)
+    # Infinite or NaN results occur when outside the support of α or β, 
+    # and also when one measure is singular wrt the other. Computing the base
+    # measures first is often much cheaper, and allows the numerically-intensive
+    # computation to "fall through" in these cases.
+    # TODO: Add tests to check that NaN cases work properly
+    ℓ = logdensity(α, β, x)
+    isfinite(ℓ) || return ℓ
 
-    return ℓμ - ℓν + logdensity(α, β, x)
+    ℓ += logdensity(μ, x)
+    ℓ -= logdensity(ν, x)
+
+    return ℓ
 end
 
 logdensity(::Lebesgue, ::Lebesgue, x) = 0.0

test/runtests.jl

@@ -154,13 +154,12 @@ end
 @testset "Half" begin
     Normal() = ∫exp(x -> -0.5x^2, Lebesgue(ℝ))
     @half Normal
-    @test logdensity(HalfNormal(), -0.2) == -Inf
+    @test logdensity(HalfNormal(), Lebesgue(ℝ), -0.2) == -Inf
     @test logdensity(HalfNormal(), 0.2) == logdensity(Normal(), 0.2)
-    
-    @half Lebesgue
-    @test basemeasure(HalfLebesgue(ℝ)) == 2 * Lebesgue(ℝ)
+    @test density(HalfNormal(), Lebesgue(ℝ), 0.2) ≈ 2 * density(Normal(), Lebesgue(ℝ), 0.2)
 end
 
+
 # import MeasureBase.:⋅
 # function ⋅(μ::Normal, kernel) 
 #     m = kernel(μ)