Smart constructors (#18)

cscherrer · web-flow · commit 85c00e5242a4 · 2021-09-07T10:01:25.000-07:00
* update 3-arg logdensity

* update logjac

* factoredbase

* update power measure

* update half.jl to use Factoredbase

* drop outdated test

* add a test

* bump version

* start on smart constructors

* work on smart constructors

* get tests passing

* roll back superpose.jl for now

* smart constructor updates

* uncomment some tests

* add missing Affine method

* update RestrictedMeasure

* Get MeasureTheory tests passing

* add fake Returns for VERSION &lt; v"1.7.0-beta1.0"

* update Breakage

* Make JET happy

* bump version
diff --git a/.github/workflows/Breakage.yml b/.github/workflows/Breakage.yml
@@ -13,9 +13,7 @@ jobs:
       fail-fast: false
       matrix:
         pkg: [
-          "cscherrer/MeasureTheory.jl",
-          "cscherrer/Soss.jl",
-          "mschauer/Mitosis.jl"
+          "cscherrer/MeasureTheory.jl"
         ]
         pkgversion: [latest, stable]
 
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "MeasureBase"
 uuid = "fa1605e6-acd5-459c-a1e6-7e635759db14"
 authors = ["Chad Scherrer <chad.scherrer@gmail.com> and contributors"]
-version = "0.4.0"
+version = "0.4.1"
 
 [deps]
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl
@@ -34,6 +34,11 @@ Methods for computing density relative to other measures will be
 """
 function logdensity end
 
+
+if VERSION < v"1.7.0-beta1.0"
+    @eval Returns(x) = _ -> x
+end
+
 include("combinators/half.jl")
 include("exp.jl")
 include("domains.jl")
@@ -60,6 +65,8 @@ include("combinators/spikemixture.jl")
 include("kernel.jl")
 include("combinators/likelihood.jl")
 include("combinators/pointwise.jl")
+include("combinators/restricted.jl")
+include("combinators/smart-constructors.jl")
 
 include("rand.jl")
 
diff --git a/src/combinators/affine.jl b/src/combinators/affine.jl
@@ -34,13 +34,11 @@ logjac(f::AffineTransform{(:μ,)}) = 0.0
 struct Affine{N,M,T} <: AbstractMeasure
     f::AffineTransform{N,T}
     parent::M
+end
 
-    function Affine(f::AffineTransform, parent::WeightedMeasure)
-        WeightedMeasure(parent.logweight, Affine(f, parent.base))
-    end
+Affine(nt::NamedTuple, μ::AbstractMeasure) = affine(nt, μ)
 
-    Affine(f::AffineTransform{N,T}, parent::M) where {N,M,T} = new{N,M,T}(f, parent)
-end
+Affine(nt::NamedTuple) = affine(nt)
 
 parent(d::Affine) = getfield(d, :parent)
 
@@ -54,10 +52,6 @@ function paramnames(::Type{A}) where {N,M, A<:Affine{N,M}}
     tuple(union(N, paramnames(M))...)
 end
 
-Affine(nt::NamedTuple, μ::AbstractMeasure) = Affine(AffineTransform(nt), μ)
-
-Affine(nt::NamedTuple) = μ -> Affine(nt, μ)
-
 Base.propertynames(d::Affine{N}) where {N} = N ∪ (:parent,)
 
 @inline function Base.getproperty(d::Affine, s::Symbol) 
@@ -77,11 +71,11 @@ logdensity(d::Affine{(:σ,)}, x) = logdensity(d.parent, d.σ \ x)
 logdensity(d::Affine{(:ω,)}, x) = logdensity(d.parent, d.ω * x)
 logdensity(d::Affine{(:μ,)}, x) = logdensity(d.parent, x - d.μ) 
 
-basemeasure(d::Affine) = Affine(getfield(d, :f), basemeasure(d.parent))
+basemeasure(d::Affine) = affine(getfield(d, :f), basemeasure(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)
+basemeasure(d::Affine{N,L}) where {N, L<:Lebesgue} = weightedmeasure(-logjac(d), d.parent)
 
 logjac(d::Affine) = logjac(getfield(d, :f))
 
diff --git a/src/combinators/factoredbase.jl b/src/combinators/factoredbase.jl
@@ -9,7 +9,7 @@ end
 
 function logdensity(d::FactoredBase, x)
     d.inbounds(x) || return -Inf
-    d.constℓ + d.varℓ
+    d.constℓ + d.varℓ()
 end
 
 basemeasure(d::FactoredBase) = d.base
diff --git a/src/combinators/half.jl b/src/combinators/half.jl
@@ -14,7 +14,7 @@ unhalf(μ::Half) = μ.parent
 function basemeasure(μ::Half)
     inbounds(x) = x > 0
     constℓ = logtwo
-    varℓ = 0.0
+    varℓ() = 0.0
     base = basemeasure(unhalf(μ))
     FactoredBase(inbounds, constℓ, varℓ, base)
 end
diff --git a/src/combinators/pointwise.jl b/src/combinators/pointwise.jl
@@ -2,9 +2,6 @@ export ⊙
 
 @concrete terse struct PointwiseProductMeasure{T} <: AbstractMeasure
     data :: T
-
-    PointwiseProductMeasure(μs...) = new{typeof(μs)}(μs)
-    PointwiseProductMeasure(μs) = new{typeof(μs)}(μs)
 end
 
 Base.size(μ::PointwiseProductMeasure) = size(μ.data)
@@ -28,30 +25,7 @@ end
 
 Base.length(m::PointwiseProductMeasure{T}) where {T} = length(m.data)
 
-function ⊙(μ::PointwiseProductMeasure{X}, ν::PointwiseProductMeasure{Y}) where {X,Y}
-    data = (μ.data..., ν.data...)
-    PointwiseProductMeasure(data...)
-end
-
-function ⊙(μ::AbstractMeasure, ν::PointwiseProductMeasure)
-    data = (μ, ν.data...)
-    PointwiseProductMeasure(data...)
-end
-
-function ⊙(μ::PointwiseProductMeasure, ν::N) where {N <: AbstractMeasure}
-    data = (μ.data..., ν)
-    PointwiseProductMeasure(data...)
-end
-
-function ⊙(μ::M, ν::N) where {M <: AbstractMeasure, N <: AbstractMeasure}
-    data = (μ, ν)
-    PointwiseProductMeasure(data...)
-end
-
-function ⊙(μ::AbstractMeasure, ℓ::Likelihood)
-    data = (μ, ℓ)
-    PointwiseProductMeasure(data...)
-end
+⊙(args...) = pointwiseproduct(args...)
 
 function logdensity(d::PointwiseProductMeasure, x)
     sum((logdensity(dⱼ, x) for dⱼ in d.data))
diff --git a/src/combinators/power.jl b/src/combinators/power.jl
@@ -32,11 +32,11 @@ export PowerMeasure
 const PowerMeasure{F,T,N,A} = ProductMeasure{F,Fill{T,N,A}}
 
 function Base.:^(μ::AbstractMeasure, dims::Integer...)
-    return μ^dims
+    return μ ^ dims
 end
 
 function Base.:^(μ::M, dims::NTuple{N,I}) where {M<:AbstractMeasure,N,I<:Integer}
-    ProductMeasure(identity, Fill(μ, dims))
+    powermeasure(μ, dims)
 end
 
 # Same as PowerMeasure
@@ -53,31 +53,32 @@ end
 
 # sampletype(d::PowerMeasure{M,N}) where {M,N} = @inbounds Array{sampletype(first(marginals(d))), N}
 
-function Base.:^(μ::WeightedMeasure, dims::NTuple{N,I}) where {N,I<:Integer}
-    k = prod(dims) * μ.logweight
-    return WeightedMeasure(k, μ.base^dims)
-end
+
 
 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)))))
+# Same as PowerMeasure
+@inline basemeasure(d::ProductMeasure{F,<:Fill}) where {F}= _basemeasure(d, (basemeasure(d.f(first(d.pars)))))
 
-@inline _basemeasure(d::PowerMeasure, b) = b ^ size(d.pars)
+# Same as PowerMeasure
+@inline _basemeasure(d::ProductMeasure{F,<:Fill}, b) where {F} = b ^ size(d.pars)
 
-@inline function _basemeasure(d::PowerMeasure, b::FactoredBase)
+# Same as PowerMeasure
+@inline function _basemeasure(d::ProductMeasure{F,<:Fill}, b::FactoredBase) where {F}
     n = length(d.pars)
     inbounds(x) = all(xj -> b.inbounds(xj), x)
     constℓ = n * b.constℓ
-    varℓ = n * b.varℓ
+    varℓ() = n * b.varℓ()
     base = b.base ^ size(d.pars)
     FactoredBase(inbounds, constℓ, varℓ, base)
 end
 
-function logdensity(d::PowerMeasure, x)
+# Same as PowerMeasure
+function logdensity(d::ProductMeasure{F,<:Fill}, x) where {F}
     d1 = d.f(first(d.pars))
     sum(xj -> logdensity(d1, xj), x)
 end
diff --git a/src/combinators/product.jl b/src/combinators/product.jl
@@ -8,16 +8,13 @@ struct ProductMeasure{F,I} <: AbstractMeasure
     pars::I
 end
 
-ProductMeasure(nt::NamedTuple) = ProductMeasure(identity, nt)
-
 Base.size(μ::ProductMeasure) = size(marginals(μ))
 
 Base.length(m::ProductMeasure{T}) where {T} = length(marginals(μ))
 
 # TODO: Pull weights outside
 basemeasure(d::ProductMeasure) = ProductMeasure(basemeasure ∘ d.f, d.pars)
 
-
 export marginals
 
 function marginals(d::ProductMeasure{F,I}) where {F,I}
@@ -39,9 +36,6 @@ function Base.show(io::IO, μ::ProductMeasure{NamedTuple{N,T}}) where {N,T}
     print(io, "Product(",μ.data, ")")
 end
 
-
-
-
 function Base.show_unquoted(io::IO, μ::ProductMeasure, indent::Int, prec::Int)
     io = IOContext(io, :compact => true)
     if Base.operator_precedence(:*) ≤ prec
@@ -141,11 +135,6 @@ end
     return x
 end
 
-
-
-
-
-
 export rand!
 using Random: rand!, GLOBAL_RNG, AbstractRNG
 
diff --git a/src/combinators/restricted.jl b/src/combinators/restricted.jl
@@ -0,0 +1,14 @@
+struct RestrictedMeasure{F,M} <: AbstractMeasure
+    f::F
+    base::M
+end
+
+function logdensity(d::RestrictedMeasure, x)
+    d.f(x) || return -Inf
+end
+
+function density(d::RestrictedMeasure, x)
+    d.f(x) || return 0.0
+end
+
+basemeasure(μ::RestrictedMeasure) = μ.base
diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl
@@ -0,0 +1,100 @@
+
+###############################################################################
+# Affine
+
+affine(f::AffineTransform, μ::AbstractMeasure) = Affine(f, μ)
+
+affine(nt::NamedTuple, μ::AbstractMeasure) = affine(AffineTransform(nt), μ)
+
+affine(f) = μ -> affine(f, μ)
+
+function affine(f::AffineTransform, parent::WeightedMeasure)
+    WeightedMeasure(parent.logweight, affine(f, parent.base))
+end
+
+function affine(f::AffineTransform, parent::FactoredBase)
+    constℓ = parent.constℓ
+    varℓ = parent.varℓ
+    # Avoid transforming `inbounds`, which is expensive
+    base = affine(f, restrict(parent.inbounds, parent.base))
+    FactoredBase(Returns(true), constℓ, varℓ, base)
+end
+
+
+###############################################################################
+# Half
+
+half(μ::AbstractMeasure) = Half(μ)
+
+###############################################################################
+# PointwiseProductMeasure
+
+pointwiseproduct(μ::AbstractMeasure...) = PointwiseProductMeasure(μ)
+
+function pointwiseproduct(μ::PointwiseProductMeasure{X}, ν::PointwiseProductMeasure{Y}) where {X,Y}
+    data = (μ.data..., ν.data...)
+    pointwiseproduct(data...)
+end
+
+function pointwiseproduct(μ::AbstractMeasure, ν::PointwiseProductMeasure)
+    data = (μ, ν.data...)
+    return pointwiseproduct(data...)
+end
+
+function pointwiseproduct(μ::PointwiseProductMeasure, ν::N) where {N <: AbstractMeasure}
+    data = (μ.data..., ν)
+    return pointwiseproduct(data...)
+end
+
+function pointwiseproduct(μ::AbstractMeasure, ℓ::Likelihood)
+    data = (μ, ℓ)
+    return PointwiseProductMeasure(data)
+end
+
+###############################################################################
+# PowerMeaure
+
+function powermeasure(μ::M, dims::NTuple{N,I}) where {M<:AbstractMeasure,N,I<:Integer}
+    productmeasure(identity, Fill(μ, dims))
+end
+
+function powermeasure(μ::WeightedMeasure, dims::NTuple{N,I}) where {N,I<:Integer}
+    k = prod(dims) * μ.logweight
+    return weightedmeasure(k, μ.base^dims)
+end
+
+###############################################################################
+# ProductMeasure
+
+productmeasure(f, pars) = ProductMeasure(f, pars)
+
+productmeasure(nt::NamedTuple) = productmeasure(identity, nt)
+
+###############################################################################
+# RestrictedMeasure
+
+restrict(f, b) = RestrictedMeasure(f, b)
+
+###############################################################################
+# SuperpositionMeasure
+
+superpose(a::AbstractArray) = SuperpositionMeasure(a)
+
+superpose(t::Tuple) = SuperpositionMeasure(t)
+superpose(nt::NamedTuple) = SuperpositionMeasure(nt)
+
+function superpose(μ::AbstractMeasure, ν::AbstractMeasure)
+    components = (μ, ν)
+    superpose(components)
+end
+
+###############################################################################
+# WeightedMeasure
+
+function weightedmeasure(ℓ::R, b::M) where {R,M}
+    WeightedMeasure{R,M}(ℓ, b)
+end 
+
+function weightedmeasure(ℓ, b::WeightedMeasure)
+    weightedmeasure(ℓ + b.logweight, b.base)
+end 
diff --git a/src/combinators/spikemixture.jl b/src/combinators/spikemixture.jl
@@ -1,5 +1,7 @@
 export SpikeMixture
 
+# TODO: Add `AbstractSuperposition <: AbstractMeasure`, and make SpikeMixture a
+# subtype of this 
 struct SpikeMixture{T,S} <: AbstractMeasure
     m::T # parent
     w::S # relative weight of parent
diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl
@@ -4,12 +4,9 @@ export SuperpositionMeasure
     struct SuperpositionMeasure{X,NT} <: AbstractMeasure
         components :: NT
     end
-
 Superposition of measures is analogous to mixture distributions, but (because
 measures need not be normalized) requires no scaling.
-
 The superposition of two measures μ and ν can be more concisely written as μ + ν.
-
 Superposition measures satisfy
     
     basemeasure(μ + ν) == basemeasure(μ) + basemeasure(ν)
@@ -43,12 +40,12 @@ testvalue(μ::SuperpositionMeasure) = testvalue(first(μ.components))
 
 function Base.:+(μ::AbstractMeasure, ν::AbstractMeasure)
     components = (μ, ν)
-    SuperpositionMeasure(components)
+    superpose(components)
 end
 
 logdensity(μ::SuperpositionMeasure, x) = logsumexp((logdensity(m,x) for m in μ.components))
 
-basemeasure(μ::SuperpositionMeasure) = SuperpositionMeasure(map(basemeasure, μ.components))
+basemeasure(μ::SuperpositionMeasure) = superpose(map(basemeasure, μ.components))
 
 # TODO: Fix `rand` method (this one is wrong)
 # function Base.rand(μ::SuperpositionMeasure{X,N}) where {X,N}
diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl
diff --git a/src/macros.jl b/src/macros.jl
diff --git a/test/runtests.jl b/test/runtests.jl
