Some convenience functions for WeightedMeasures (#43)

* convenience functions for WeightedMeasures

* Fix typo

* Second typo

* Add tests and use _logweight

* Update test/combinators/weighted.jl

Co-authored-by: Chad Scherrer <[email protected]>

* Remove length

Co-authored-by: Chad Scherrer <[email protected]>

Author: theogf

src/combinators/smart-constructors.jl

@@ -62,8 +62,8 @@ function productmeasure(f::Returns{FB}, param_maps, pars) where {FB<:FactoredBas
 end
 
 function productmeasure(f::Returns{W}, ::typeof(identity), pars) where {W<:WeightedMeasure}
-    ℓ = f.value.logweight
-    base = f.value.base
+    ℓ = _logweight(f.value)
+    base = basemeasure(f.value)
     newbase = productmeasure(Returns(base), identity, pars)
     weightedmeasure(length(pars) * ℓ, newbase)
 end
@@ -102,7 +102,7 @@ function weightedmeasure(ℓ::R, b::M) where {R,M}
 end
 
 function weightedmeasure(ℓ, b::WeightedMeasure)
-    weightedmeasure(ℓ + b.logweight, b.base)
+    weightedmeasure(ℓ + _logweight(b), b.base)
 end
 
 ###############################################################################

src/combinators/weighted.jl

@@ -9,20 +9,27 @@ export WeightedMeasure, AbstractWeightedMeasure
 
 abstract type AbstractWeightedMeasure <: AbstractMeasure end
 
-logweight(μ::AbstractWeightedMeasure) = μ.logweight
-basemeasure(μ::AbstractWeightedMeasure) = μ.base
+# By default the weight for all measure is 1
+_logweight(::AbstractMeasure) = 0
 
-@inline function logdensity_def(d::AbstractWeightedMeasure, x)
+@inline function logdensity_def(d::AbstractWeightedMeasure, _)
     d.logweight
 end
 
+function Base.rand(rng::AbstractRNG, ::Type{T}, μ::AbstractWeightedMeasure) where {T}
+    rand(rng, T, basemeasure(μ))
+end
+
 ###############################################################################
 
 struct WeightedMeasure{R,M} <: AbstractWeightedMeasure
     logweight::R
     base::M
 end
 
+_logweight(μ::WeightedMeasure) = μ.logweight
+basemeasure(μ::AbstractWeightedMeasure) = μ.base
+
 function Base.show(io::IO, μ::WeightedMeasure)
     io = IOContext(io, :compact => true)
     print(io, exp(μ.logweight), " * ", μ.base)

test/combinators/weighted.jl

@@ -0,0 +1,20 @@
+using Random: MersenneTwister
+using Test
+
+using MeasureBase
+using MeasureBase: _logweight, weightedmeasure, WeightedMeasure
+
+@testset "weighted" begin
+    @test iszero(_logweight(Lebesgue(ℝ)))
+    μ₀ = Dirac(0.0)
+    w = 2.0
+    μ = @inferred w * μ₀
+    @test μ == WeightedMeasure(log(w), μ₀) == weightedmeasure(log(w), μ₀)
+    @test μ isa WeightedMeasure
+    @test _logweight(μ) == log(w)
+    @test _logweight(w * μ) == 2 * log(w)
+    @test rand(MersenneTwister(123), μ) == rand(MersenneTwister(123), μ₀) 
+    x = rand()
+    @test logdensity_def(μ, x) == log(w)
+    @test logdensityof(μ, x) == logdensityof(μ₀, x)
+end

test/runtests.jl

@@ -212,3 +212,5 @@ end
 #         @test f(x) ≈ x^2
 #     end
 end
+
+include("combinators/weighted.jl")