cleanup (#39)

- Get `JET.report_package(MeasureBase)` to pass
- Drop the `partialstatic` stuff, fun idea but didn't seem to help much
- Add dependencies for `LogarithmicNumbers` and `StatsFuns`
- More docstrings
- Move `ConditionalMeasure` (`|`) here from MeasureTheory 
- Improvements for `superpose`
- Some minor tidying up

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.6.0"
+version = "0.6.1"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
@@ -11,12 +11,14 @@ FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
+LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899"
 MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
 PrettyPrinting = "54e16d92-306c-5ea0-a30b-337be88ac337"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 
@@ -27,10 +29,12 @@ DensityInterface = "0.4"
 FillArrays = "0.12, 0.13"
 IfElse = "0.1"
 LogExpFunctions = "0.3"
+LogarithmicNumbers = "1"
 MappedArrays = "0.4"
 PrettyPrinting = "0.3, 0.4"
 Reexport = "1"
 Static = "0.5, 0.6"
+StatsFuns = "0.9"
 Tricks = "0.1"
 julia = "1.3"
 

src/MeasureBase.jl

@@ -28,6 +28,25 @@ export rebase
 
 export AbstractMeasure
 
+import IfElse: ifelse
+export logdensity_def
+export basemeasure
+export basekleisli
+
+"""
+    inssupport(m, x)
+    insupport(m)
+
+`insupport(m,x)` computes whether `x` is in the support of `m`.
+
+`insupport(m)` returns a function, and satisfies
+
+    insupport(m)(x) == insupport(m, x)
+"""
+function insupport end
+
+export insupport
+
 abstract type AbstractMeasure end
 
 using Static: @constprop
@@ -45,28 +64,35 @@ gentype(μ::AbstractMeasure) = typeof(testvalue(μ))
 
 # gentype(μ::AbstractMeasure) = gentype(basemeasure(μ))
 
-import IfElse: ifelse
-export logdensity_def
-export basemeasure
-export basekleisli
-
 using LogExpFunctions: logsumexp
 
 """
-    logdensity_def(μ::AbstractMeasure{X}, x::X)
+`logdensity_def` is the standard way to define a log-density for a new measure.
+Note that this definition does not include checking for membership in the
+support; this is instead checked using `insupport`. `logdensity_def` is
+a low-level function, and should typically not be called directly. See
+`logdensityof` for more information and other alternatives.
+
+---
+
+    logdensity_def(m, x)
 
-Compute the logdensity of the measure μ at the point x. This is the standard way
-to define `logdensity` for a new measure. the base measure is implicit here, and
-is understood to be `basemeasure(μ)`.
+Compute the log-density of the measure m at the point `x`, relative to
+`basemeasure(m)`, and assuming `insupport(m, x)`.
 
-Methods for computing density relative to other measures will be
+---
+
+    logdensity_def(m1, m2, x)
+
+Compute the log-density of `m1` relative to `m2` at the point `x`, assuming
+`insupport(m1, x)` and `insupport(m2, x)`.
 """
 function logdensity_def end
 
 using Compat
 
+include("schema.jl")
 include("splat.jl")
-include("partial-static.jl")
 include("proxies.jl")
 include("kleisli.jl")
 include("parameterized.jl")
@@ -81,6 +107,7 @@ include("primitives/lebesgue.jl")
 include("primitives/dirac.jl")
 include("primitives/trivial.jl")
 
+include("combinators/conditional.jl")
 include("combinators/bind.jl")
 include("combinators/transformedmeasure.jl")
 include("combinators/factoredbase.jl")
@@ -97,21 +124,6 @@ include("combinators/smart-constructors.jl")
 include("rand.jl")
 
 include("density.jl")
-module Interface
-
-using Reexport
-using MeasureBase
-using MeasureBase:basemeasure_depth, proxy
-using MeasureBase: insupport
-@reexport using Test
-
-export test_interface
-export basemeasure_depth
-export proxy
-export insupport
-
-include("interface.jl")
-end # module Interface
 
 using .Interface
 

src/combinators/conditional.jl

@@ -0,0 +1,28 @@
+struct ConditionalMeasure{M,C} <: AbstractMeasure
+    parent::M 
+    constraint::C
+end
+
+"""
+    (m::AbstractMeasure) | constraint
+
+Return a new measure by constraining `m` to satisfy `constraint`.
+
+Note that the form of `constraint` will vary depending on the structure of a
+given measure. For example, a measure over `NamedTuple`s may allow `NamedTuple`
+constraints, while another may require `constraint` to be a predicate or a
+function returning a real number (in which case the constraint could be
+considered as the zero-set of that function). 
+
+At the time of this writing, invariants required of this function are not yet
+settled. Specifically, there's the question of normalization. It's common for
+conditional distributions to be normalized, but this can often not be expressed
+in closed form, and can be very expensive to compute. For more general measures,
+the notion of normalization may not even make sense.
+
+Because of this, this interface is not yet stable, and users should expect
+upcoming changes.
+"""
+Base.:|(μ::AbstractMeasure, constraint) = ConditionalMeasure(μ, constraint)
+
+@inline basemeasure(cm::ConditionalMeasure) = basemeasure(cm.parent) | cm.constraint

src/combinators/half.jl

@@ -15,7 +15,7 @@ isnonnegative(x) = x ≥ 0.0
 
 @inline function basemeasure(μ::Half)
     constℓ = static(logtwo)
-    varℓ = Returns(static(0.0))
+    varℓ = Returns(0.0)
     base = basemeasure(unhalf(μ))
     FactoredBase(constℓ, varℓ, base)
 end

src/combinators/product.jl

@@ -100,13 +100,13 @@ end
     end
 
     vals = map(x -> Expr(:(=), x,x), N)
-    push!(q, Expr(:tuple, vals...))
+    push!(q.args, Expr(:tuple, vals...))
     return q
 end
 
 function basemeasure(μ::ProductMeasure{Base.Generator{I,F}}) where {I,F}
     mar = marginals(μ)
-    T = Core.Compiler.return_type(mar.f, Tuple{_eltype(mar.iter)})
+    T = Core.Compiler.return_type(mar.f, Tuple{eltype(mar.iter)})
     B = Core.Compiler.return_type(basemeasure, Tuple{T})
     _basemeasure(μ, B, static(Base.issingletontype(B)))
 end

src/combinators/smart-constructors.jl

@@ -76,6 +76,14 @@ superpose(a::AbstractArray) = SuperpositionMeasure(a)
 superpose(t::Tuple) = SuperpositionMeasure(t)
 superpose(nt::NamedTuple) = SuperpositionMeasure(nt)
 
+function superpose(μ::T, ν::T) where {T}
+    if μ==ν
+        return weightedmeasure(static(logtwo), μ)
+    else
+        return superpose((μ, ν))
+    end
+end
+
 function superpose(μ::AbstractMeasure, ν::AbstractMeasure)
     components = (μ, ν)
     superpose(components)

src/combinators/superpose.jl

@@ -58,10 +58,53 @@ function Base.:+(μ::AbstractMeasure, ν::AbstractMeasure)
     superpose(components)
 end
 
-function logdensity_def(μ::SuperpositionMeasure, x)
-    logsumexp((logdensity_def(m, x) for m in μ.components))
+using LogarithmicNumbers
+
+oneplus(x::ULogarithmic) = exp(ULogarithmic, log1pexp(x.log))
+
+@inline function density_def(s::SuperpositionMeasure{Tuple{A,B}}, x) where {A,B}
+    (μ, ν) = s.components
+    insupport(μ, x) || return exp(ULogarithmic, logdensity_def(ν, x))
+    insupport(ν, x) || return exp(ULogarithmic, logdensity_def(μ, x))
+    α = basemeasure(μ)
+    β = basemeasure(ν)
+    dμ_dα = exp(ULogarithmic, logdensity_def(μ, x))
+    dν_dβ = exp(ULogarithmic, logdensity_def(ν, x))
+    dα_dβ = exp(ULogarithmic, logdensity_rel(α, β, x))
+    dβ_dα = inv(dα_dβ)
+    return dμ_dα / oneplus(dβ_dα) + dν_dβ / oneplus(dα_dβ)
 end
 
+using StatsFuns
+
+@inline function logdensity_def(μ::T, ν::T, x::Any) where T<:(SuperpositionMeasure{Tuple{A, B}} where {A, B})
+    if μ === ν
+        return zero(return_type(logdensity_def, (μ, x)))
+    else
+        return logdensity_def(μ,x) - logdensity_def(ν, x)
+    end
+end
+
+@inline function logdensity_def(s::SuperpositionMeasure{Tuple{A,B}}, β, x) where {A,B}
+    (μ, ν) = s.components
+    insupport(μ, x) || return logdensity_rel(ν, β, x)
+    insupport(ν, x) || return logdensity_rel(μ, β, x)
+    return logaddexp(logdensity_rel(μ, β, x), logdensity_rel(ν, β, x))
+end
+
+@inline function logdensity_def(s::SuperpositionMeasure{Tuple{A,B}}, β::SuperpositionMeasure, x) where {A,B}
+    (μ, ν) = s.components
+    insupport(μ, x) || return logdensity_rel(ν, β, x)
+    insupport(ν, x) || return logdensity_rel(μ, β, x)
+    return logaddexp(logdensity_rel(μ, β, x), logdensity_rel(ν, β, x))
+end
+
+@inline function logdensity_def(s, β::SuperpositionMeasure{Tuple{A,B}}, x) where {A,B}
+    -logdensity_def(β, s, x)
+end
+
+@inline logdensity_def(s::SuperpositionMeasure, x) = log(density_def(s, x))
+
 basemeasure(μ::SuperpositionMeasure) = superpose(map(basemeasure, μ.components))
 
 # TODO: Fix `rand` method (this one is wrong)