Affine (#23)

# `basekernel`

For `k::Kernel`, `basekernel` has the property
```julia
basemeasure(k(p)) == basekernel(k)(p)
```
There are many cases for which
```julia
basekernel(k) isa Returns
```
in which case `p` doesn't matter. This can make things much more efficient, especially for large product measures.

# Smart Constructors

Many functions are now in terms of "smart constructors". These have lower-case names, for example `kernel` as opposed to `Kernel`. There's also a new `TupleProductMeasure` to handle cases like `Normal() ⊗ Lebesgue(ℝ)`.



A given smart constructor should have a very small number of methods that call the struct itself; the remainder should all make a small reduction step, usually calling another smart constructor.

# Changes to ProductMeasure

`ProductMeasure` is now in terms of `Kernel`.

# MapsTo

When computing the logpdf of an affine measure, we currently have to recompute the pullback many times. As a step toward improving this, this PR introduces a `MapsTo` combinator that works like 
```julia
julia> typeof([1,2] ↦ 3)
MeasureBase.MapsTo{Vector{Int64}, Int64}
```
The idea is that we can carry this around until the logpdf computation is done, and avoid recomputation in this way. 

-----------------------------------

* WIP

* have logpdf use == instead of === (at least for now)

* move kernel stuff to smart-constructors

* some refactoring

* fixing stuff

* typo

* AbstractProductMeasure

* get tests passing

* shorten basemeasure(d::ProductMeasure)

* basekernel docstring

* correct dispatch problem

* update comment

* thunks WIP

* 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.4.4"
+version = "0.4.5"
 
 [deps]
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"

src/MeasureBase.jl

@@ -5,6 +5,7 @@ const logtwo = log(2.0)
 using Random
 import Random: rand!
 
+using FillArrays
 using ConcreteStructs
 using MLStyle
 
@@ -21,6 +22,7 @@ sampletype(μ::AbstractMeasure) = typeof(testvalue(μ))
 
 export logdensity
 export basemeasure
+export basekernel
 
 using LogExpFunctions: logsumexp
 
@@ -46,12 +48,14 @@ if VERSION < v"1.7.0-beta1.0"
     end
 end
 
+include("kernel.jl")
+include("parameterized.jl")
+include("combinators/mapsto.jl")
 include("combinators/half.jl")
 include("exp.jl")
 include("domains.jl")
 include("utils.jl")
 include("absolutecontinuity.jl")
-include("parameterized.jl")
 include("macros.jl")
 include("resettablerng.jl")
 
@@ -69,7 +73,6 @@ include("combinators/for.jl")
 include("combinators/power.jl")
 include("combinators/affine.jl")
 include("combinators/spikemixture.jl")
-include("kernel.jl")
 include("combinators/likelihood.jl")
 include("combinators/pointwise.jl")
 include("combinators/restricted.jl")

src/combinators/affine.jl

@@ -34,6 +34,12 @@ Base.size(f::AffineTransform, n::Int) = @inbounds size(f)[n]
 (f::AffineTransform{(:ω,)})(x) = f.ω \ x
 (f::AffineTransform{(:μ,:σ)})(x) = f.σ * x + f.μ
 (f::AffineTransform{(:μ,:ω)})(x) = f.ω \ x + f.μ
+1
+rowsize(x) = ()
+rowsize(x::AbstractArray) = (size(x,1),)
+
+colsize(x) = ()
+colsize(x::AbstractArray) = (size(x,2),)
 
 @inline function apply!(x, f::AffineTransform{(:μ,)}, z)
     x .= z .+ f.μ
@@ -130,14 +136,37 @@ Base.size(d::Affine) = size(d.μ)
 Base.size(d::Affine{(:σ,)}) = (size(d.σ, 1),)
 Base.size(d::Affine{(:ω,)}) = (size(d.ω, 2),)
 
-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.μ) 
-logdensity(d::Affine{(:μ,:σ)}, x) = logdensity(d.parent, d.σ \ (x - d.μ)) 
-logdensity(d::Affine{(:μ,:ω)}, x) = logdensity(d.parent, d.ω * (x - d.μ)) 
+logdensity(d::Affine, x::MapsTo) = logdensity(d.parent, x.x)
+
+function logdensity(d::Affine{(:σ,)}, x)
+    z = d.σ \ x
+    # @show z
+    # println()
+    logdensity(d.parent, z)
+end
+
+function logdensity(d::Affine{(:ω,)}, x)
+    z = d.ω * x
+    logdensity(d.parent, z)
+end
+
+function logdensity(d::Affine{(:μ,)}, x)
+    z = x - d.μ
+    logdensity(d.parent, z) 
+end
+
+function logdensity(d::Affine{(:μ,:σ)}, x)
+    z = d.σ \ (x - d.μ)
+    logdensity(d.parent, z) 
+end
+
+function logdensity(d::Affine{(:μ,:ω)}, x)
+    z = d.ω * (x - d.μ)
+    logdensity(d.parent, z)
+end 
 
 # logdensity(d::Affine{(:μ,:ω)}, x) = logdensity(d.parent, d.σ \ (x - d.μ))
-@inline function logdensity(d::Affine{(:μ,:σ), Tuple{AbstractVector, AbstractMatrix}}, x)
+@inline function logdensity(d::Affine{(:μ,:σ), P, Tuple{V,M}}, x) where {P, V<:AbstractVector, M<:AbstractMatrix}
     z = x - d.μ
     σ = d.σ
     if σ isa Factorization
@@ -149,7 +178,7 @@ logdensity(d::Affine{(:μ,:ω)}, x) = logdensity(d.parent, d.ω * (x - d.μ))
 end
 
 # logdensity(d::Affine{(:μ,:ω)}, x) = logdensity(d.parent, d.ω * (x - d.μ))
-@inline function logdensity(d::Affine{(:μ,:ω), Tuple{AbstractVector, AbstractMatrix}}, x)
+@inline function logdensity(d::Affine{(:μ,:ω), P,Tuple{V,M}}, x) where {P,V<:AbstractVector, M<:AbstractMatrix}
     z = x - d.μ
     lmul!(d.ω, z)
     logdensity(d.parent, z)
@@ -161,7 +190,7 @@ basemeasure(d::Affine) = affine(getfield(d, :f), basemeasure(d.parent))
 # 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)
 
-function basemeasure(d::Affine{N,L}) where {N, L<:PowerMeasure{typeof(identity), <:Lebesgue}}
+function basemeasure(d::Affine{N,L}) where {N, L<:PowerMeasure{typeof(identity), typeof(identity), <:Lebesgue}}
     weightedmeasure(-logjac(d), d.parent)
 end
 

src/combinators/for.jl

@@ -75,12 +75,12 @@ julia> For(eachrow(rand(4,2))) do x Normal(x[1], x[2]) end |> marginals |> colle
 ```
 
 """
-For(f, dims...) = ProductMeasure(i -> f(i...), zip(dims...))
+For(f, dims...) = productmeasure(i -> f(i...), zip(dims...))
 
-For(f, inds::AbstractArray) = ProductMeasure(f, inds)
+For(f, inds::AbstractArray) = productmeasure(f, inds)
 
-For(f, n::Int) = ProductMeasure(f, 1:n)
-For(f, dims::Int...) = ProductMeasure(i -> f(Tuple(i)...), CartesianIndices(dims))
+For(f, n::Int) = productmeasure(f, 1:n)
+For(f, dims::Int...) = productmeasure(i -> f(Tuple(i)...), CartesianIndices(dims))
 
 
 function Base.eltype(d::ProductMeasure{F,I}) where {F,I<:AbstractArray}

src/combinators/mapsto.jl

@@ -0,0 +1,19 @@
+struct MapsTo{X,Y}
+    x::X
+    y::Y
+end
+
+export ↦, mapsto
+
+mapsto(x,y) = x ↦ y
+
+↦(x::X,y::Y) where {X,Y} = MapsTo{X,Y}(x,y)
+
+Base.first(t::MapsTo) = t.x
+Base.last(t::MapsTo) = t.y
+
+Base.Pair(t::MapsTo) = t.x => t.y
+
+Base.show(io::IO, t::MapsTo) = print(t.x, " ↦ ", t.y)
+
+logdensity(d, t::MapsTo) = logdensity(d, t.y)

src/combinators/power.jl

@@ -29,7 +29,9 @@ using FillArrays: Fill
 
 export PowerMeasure
 
-const PowerMeasure{F,T,N,A} = ProductMeasure{F,Fill{T,N,A}}
+const PowerMeasure{F,S,T,N,A} = ProductMeasure{F,S,Fill{T,N,A}}
+
+Base.:^(μ::AbstractMeasure, ::Tuple{}) = μ
 
 function Base.:^(μ::AbstractMeasure, dims::Integer...)
     return μ ^ dims
@@ -40,13 +42,13 @@ function Base.:^(μ::M, dims::NTuple{N,I}) where {M<:AbstractMeasure,N,I<:Intege
 end
 
 # Same as PowerMeasure
-function Base.show(io::IO, d::ProductMeasure{F,<:Fill}) where {F}
+function Base.show(io::IO, d::ProductMeasure{F,S,<:Fill}) where {F,S}
     io = IOContext(io, :compact => true)
     print(io, d.f(first(d.pars)), " ^ ", size(d.pars))
 end
 
 # Same as PowerMeasure{F,T,1} where {F,T}
-function Base.show(io::IO, d::ProductMeasure{F,Fill{T,1,A}}) where {F,T,A}
+function Base.show(io::IO, d::ProductMeasure{F,S,Fill{T,1,A}}) where {F,S,T,A}
     io = IOContext(io, :compact => true)
     print(io, d.f(first(d.pars)), " ^ ", size(d.pars)[1])
 end
@@ -55,20 +57,24 @@ end
 
 
 
-params(d::ProductMeasure{F,<:Fill}) where {F} = params(first(marginals(d)))
+params(d::ProductMeasure{F,S,<:Fill}) where {F,S} = params(first(marginals(d)))
 
-params(::Type{P}) where {F,P<:ProductMeasure{F,<:Fill}} = params(D)
+params(::Type{P}) where {F,S,P<:ProductMeasure{F,S,<:Fill}} = params(D)
 
 # basemeasure(μ::PowerMeasure) = @inbounds basemeasure(first(μ.data))^size(μ.data)
 
 # Same as PowerMeasure
-@inline basemeasure(d::ProductMeasure{F,<:Fill}) where {F}= _basemeasure(d, (basemeasure(d.f(first(d.pars)))))
+@inline function basemeasure(d::ProductMeasure{F,S,<:Fill}) where {F,S}
+    _basemeasure(d, (basemeasure(d.f(first(d.pars)))))
+end
 
 # Same as PowerMeasure
-@inline _basemeasure(d::ProductMeasure{F,<:Fill}, b) where {F} = b ^ size(d.pars)
+@inline function _basemeasure(d::ProductMeasure{F,S,<:Fill}, b) where {F,S}
+    b ^ size(d.pars)
+end
 
 # Same as PowerMeasure
-@inline function _basemeasure(d::ProductMeasure{F,<:Fill}, b::FactoredBase) where {F}
+@inline function _basemeasure(d::ProductMeasure{F,S,<:Fill}, b::FactoredBase) where {F,S}
     n = length(d.pars)
     inbounds(x) = all(xj -> b.inbounds(xj), x)
     constℓ = n * b.constℓ
@@ -78,7 +84,7 @@ params(::Type{P}) where {F,P<:ProductMeasure{F,<:Fill}} = params(D)
 end
 
 # Same as PowerMeasure
-@inline function logdensity(d::ProductMeasure{F,<:Fill}, x) where {F}
+@inline function logdensity(d::ProductMeasure{F,S,<:Fill}, x) where {F,S}
     d1 = d.f(first(d.pars))
     sum(xj -> logdensity(d1, xj), x)
 end

src/combinators/product.jl

@@ -2,39 +2,54 @@ export ProductMeasure
 
 using MappedArrays
 using Base: @propagate_inbounds
+import Base
 using FillArrays
 
-struct ProductMeasure{F,I} <: AbstractMeasure
-    f::F
+abstract type AbstractProductMeasure <: AbstractMeasure end
+
+
+struct ProductMeasure{F,S,I} <: AbstractProductMeasure
+    f::Kernel{F,S}
     pars::I
 end
 
+
+# TODO: Test for equality without traversal, probably by first converting to a
+# canonical form
+function Base.:(==)(a::ProductMeasure, b::ProductMeasure)
+    all(zip(a.pars, b.pars)) do (aᵢ, bᵢ)
+        a.f(aᵢ) == b.f(bᵢ)
+    end
+end
+
 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)
-basemeasure(d::ProductMeasure{typeof(identity)}) = ProductMeasure(identity, map(basemeasure, d.pars))
-basemeasure(d::ProductMeasure{typeof(identity), <:FillArrays.Fill}) = ProductMeasure(identity, map(basemeasure, d.pars))
+basemeasure(d::ProductMeasure) = productmeasure(basekernel(d.f), d.pars)
+
+# TODO: Do we need these methods?
+# basemeasure(d::ProductMeasure) = ProductMeasure(basemeasure ∘ d.f, d.pars)
+# basemeasure(d::ProductMeasure{typeof(identity)}) = ProductMeasure(identity, map(basemeasure, d.pars))
+# basemeasure(d::ProductMeasure{typeof(identity), <:FillArrays.Fill}) = ProductMeasure(identity, map(basemeasure, d.pars))
 
 export marginals
 
-function marginals(d::ProductMeasure{F,I}) where {F,I}
+function marginals(d::ProductMeasure{F,S,I}) where {F,S,I}
     _marginals(d, isiterable(I))
 end
 
-function _marginals(d::ProductMeasure{F,I}, ::Iterable) where {F,I}
+function _marginals(d::ProductMeasure, ::Iterable)
     return (d.f(i) for i in d.pars)
 end
 
-function _marginals(d::ProductMeasure{F,I}, ::NonIterable) where {F,I}
+function _marginals(d::ProductMeasure{F,S,I}, ::NonIterable) where {F,S,I}
     error("Type $I is not iterable. Add an `iterate` or `marginals` method to fix.")
 end
 
 testvalue(d::ProductMeasure) = map(testvalue, marginals(d))
 
-function Base.show(io::IO, μ::ProductMeasure{NamedTuple{N,T}}) where {N,T}
+function Base.show(io::IO, μ::ProductMeasure{F,S,NamedTuple{N,T}}) where {F,S,N,T}
     io = IOContext(io, :compact => true)
     print(io, "Product(",μ.data, ")")
 end
@@ -55,21 +70,25 @@ end
 ###############################################################################
 # I <: Tuple
 
+struct TupleProductMeasure{T} <: AbstractProductMeasure
+    pars::T
+end
+
 export ⊗
-⊗(μs::AbstractMeasure...) = ProductMeasure(identity, μs)
+⊗(μs::AbstractMeasure...) = productmeasure(μs)
 
-marginals(d::ProductMeasure{F,T}) where {F, T<:Tuple} = map(d.f, d.pars)
+marginals(d::TupleProductMeasure{T}) where {F, T<:Tuple} = d.pars
 
-function Base.show(io::IO, μ::ProductMeasure{F,T}) where {F,T <: Tuple}
+function Base.show(io::IO, μ::TupleProductMeasure{T}) where {F,T <: Tuple}
     io = IOContext(io, :compact => true)
     print(io, join(string.(marginals(μ)), " ⊗ "))
 end
 
-@inline function logdensity(d::ProductMeasure{F,T}, x::Tuple) where {F,T<:Tuple}
-    mapreduce(logdensity, +, d.f.(d.pars), x)
+@inline function logdensity(d::TupleProductMeasure, x::Tuple) where {T<:Tuple}
+    mapreduce(logdensity, +, d.pars, x)
 end
 
-function Base.rand(rng::AbstractRNG, ::Type{T}, d::ProductMeasure{F,I}) where {T,F,I<:Tuple}
+function Base.rand(rng::AbstractRNG, ::Type{T}, d::TupleProductMeasure) where {T}
     rand.(d.pars)
 end
 
@@ -93,7 +112,7 @@ end
 ###############################################################################
 # I <: CartesianIndices
 
-function Base.show(io::IO, d::ProductMeasure{F,I}) where {F, I<:CartesianIndices}
+function Base.show(io::IO, d::ProductMeasure{F,S,I}) where {F, S, I<:CartesianIndices}
     io = IOContext(io, :compact => true)
     print(io, "For(")
     print(io, d.f, ", ")
@@ -102,7 +121,7 @@ function Base.show(io::IO, d::ProductMeasure{F,I}) where {F, I<:CartesianIndices
 end
 
 
-# function Base.rand(rng::AbstractRNG, ::Type{T}, d::ProductMeasure{F,I}) where {T,F,I<:CartesianIndices}
+# function Base.rand(rng::AbstractRNG, ::Type{T}, d::ProductMeasure{F,S,I}) where {T,F,I<:CartesianIndices}
 
 # end
 
@@ -114,12 +133,12 @@ export rand!
 using Random: rand!, GLOBAL_RNG, AbstractRNG
 
 
-function logdensity(d::ProductMeasure{F,I}, x) where {F, I<:Base.Generator}
+function logdensity(d::ProductMeasure{F,S,I}, x) where {F, S, I<:Base.Generator}
     sum((logdensity(dj, xj) for (dj, xj) in zip(marginals(d), x)))
 end
 
 
-function Base.rand(rng::AbstractRNG, ::Type{T}, d::ProductMeasure{F,I}) where {T,F,I<:Base.Generator}
+function Base.rand(rng::AbstractRNG, ::Type{T}, d::ProductMeasure{F,S,I}) where {T,F,S,I<:Base.Generator}
     mar = marginals(d)
     elT = typeof(rand(rng, T, first(mar)))
 
@@ -183,8 +202,8 @@ end
 #     μ.data
 # end
 
-function ConstructionBase.constructorof(::Type{P}) where {F,I,P <: ProductMeasure{F,I}}
-    p -> ProductMeasure(d.f, p)
+function ConstructionBase.constructorof(::Type{P}) where {F,S,I,P <: ProductMeasure{F,S,I}}
+    p -> productmeasure(d.f, p)
 end
 
 # function Accessors.set(d::ProductMeasure{N}, ::typeof(params), p) where {N}
@@ -201,3 +220,12 @@ end
 #         logdensity(μ_ν_x...)
 #     end
 # end
+
+function kernelfactor(μ::ProductMeasure{F,S,<:Fill}) where {F,S}
+    k = kernel(first(marginals(μ)))
+    (p -> k.f(p)^size(μ), k.ops)
+end
+
+function kernelfactor(μ::ProductMeasure{F,S,A}) where {F,S,A<:AbstractArray}
+    (p -> set.(marginals(μ), params, p), μ.pars)
+end