logjac for non-square matrices (#21)

* logjac for non-square matrices

* `rand!` for Affine measures

* testvalue(::Affine)

* check for factorization

* drop redundant method

* `apply!` method for Factorization

* bugfix

* drop `rand` for now

* roll back `rand` for now

* logdensity for product measures

* bugfix

* logpdf

* inlining

* add f to propertynames

* drop unneeded method

* affine test

* tests

* typo

* version bump

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

src/MeasureBase.jl

@@ -3,6 +3,7 @@ module MeasureBase
 const logtwo = log(2.0)
 
 using Random
+import Random: rand!
 
 using ConcreteStructs
 using MLStyle

src/combinators/affine.jl

@@ -16,14 +16,65 @@ Base.propertynames(d::AffineTransform{N}) where {N} = N
 @inline Base.inv(f::AffineTransform{(:ω,)}) = AffineTransform((σ = f.ω,))
 @inline Base.inv(f::AffineTransform{(:μ,)}) = AffineTransform((μ = -f.μ,))
 
+# `size(f) == (m,n)` means `f : ℝⁿ → ℝᵐ`  
+Base.size(f::AffineTransform{(:μ,:σ)}) = size(f.σ)
+Base.size(f::AffineTransform{(:μ,:ω)}) = size(f.ω)
+Base.size(f::AffineTransform{(:σ,)})  = size(f.σ)
+Base.size(f::AffineTransform{(:ω,)})  = size(f.ω)
+
+function Base.size(f::AffineTransform{(:μ,)})
+    (n,) = size(f.μ)
+    return (n,n)
+end
+
+Base.size(f::AffineTransform, n::Int) = @inbounds size(f)[n]
+
 (f::AffineTransform{(:μ,)})(x) = x + f.μ
 (f::AffineTransform{(:σ,)})(x) = f.σ * x
 (f::AffineTransform{(:ω,)})(x) = f.ω \ x
 (f::AffineTransform{(:μ,:σ)})(x) = f.σ * x + f.μ
 (f::AffineTransform{(:μ,:ω)})(x) = f.ω \ x + f.μ
 
+@inline function apply!(x, f::AffineTransform{(:μ,)}, z)
+    x .= z .+ f.μ
+    return x
+end
+
+@inline function apply!(x, f::AffineTransform{(:σ,)}, z)
+    mul!(x, f.σ, z)
+    return x
+end
+
+@inline function apply!(x, f::AffineTransform{(:ω,), Tuple{F}}, z) where {F<:Factorization}
+    ldiv!(x, f.ω, z)
+    return x
+end
+
+@inline function apply!(x, f::AffineTransform{(:ω,)}, z)
+    ldiv!(x, factorize(f.ω), z)
+    return x
+end
+
+@inline function apply!(x, f::AffineTransform{(:μ,:σ)}, z)
+    apply!(x, AffineTransform((σ = f.σ,)))
+    apply!(x, AffineTransform((μ = f.μ,)))
+    return x
+end
+
+@inline function apply!(x, f::AffineTransform{(:μ,:ω)}, z)
+    apply!(x, AffineTransform((ω = f.ω,)))
+    apply!(x, AffineTransform((μ = f.μ,)))
+    return x
+end
+
+function logjac(x::AbstractMatrix) 
+    (m,n) = size(x)
+    m == n && return first(logabsdet(x))
 
-logjac(x::AbstractMatrix) = first(logabsdet(x))
+    # Equivalent to sum(log, svdvals(x)), but much faster
+    m > n && return first(logabsdet(x' * x)) / 2
+    return first(logabsdet(x * x')) / 2
+end
 
 logjac(x::Number) = log(abs(x))
 
@@ -41,6 +92,12 @@ logjac(f::AffineTransform{(:μ,)}) = 0.0
     parent::M
 end
 
+function testvalue(d::Affine)
+    f = getfield(d, :f)
+    z = testvalue(parent(d))
+    return f(z)
+end
+
 Affine(nt::NamedTuple, μ::AbstractMeasure) = affine(nt, μ)
 
 Affine(nt::NamedTuple) = affine(nt)
@@ -57,17 +114,19 @@ function paramnames(::Type{A}) where {N,M, A<:Affine{N,M}}
     tuple(union(N, paramnames(M))...)
 end
 
-Base.propertynames(d::Affine{N}) where {N} = N ∪ (:parent,)
+Base.propertynames(d::Affine{N}) where {N} = N ∪ (:parent,:f)
 
 @inline function Base.getproperty(d::Affine, s::Symbol) 
     if s === :parent
         return getfield(d, :parent)
+    elseif s === :f 
+        return getfield(d, :f)
     else
         return getproperty(getfield(d, :f), s)
     end
 end
 
-Base.size(d) = size(d.μ)
+Base.size(d::Affine) = size(d.μ)
 Base.size(d::Affine{(:σ,)}) = (size(d.σ, 1),)
 Base.size(d::Affine{(:ω,)}) = (size(d.ω, 2),)
 
@@ -78,14 +137,19 @@ logdensity(d::Affine{(:μ,:σ)}, x) = logdensity(d.parent, d.σ \ (x - d.μ))
 logdensity(d::Affine{(:μ,:ω)}, x) = logdensity(d.parent, d.ω * (x - d.μ)) 
 
 # logdensity(d::Affine{(:μ,:ω)}, x) = logdensity(d.parent, d.σ \ (x - d.μ))
-function logdensity(d::Affine{(:μ,:σ), Tuple{AbstractVector, AbstractMatrix}}, x)
+@inline function logdensity(d::Affine{(:μ,:σ), Tuple{AbstractVector, AbstractMatrix}}, x)
     z = x - d.μ
-    ldiv!(d.σ, z)
+    σ = d.σ
+    if σ isa Factorization
+        ldiv!(σ, z)
+    else
+        ldiv!(factorize(σ), z)
+    end
     logdensity(d.parent, z)
 end
 
 # logdensity(d::Affine{(:μ,:ω)}, x) = logdensity(d.parent, d.ω * (x - d.μ))
-function logdensity(d::Affine{(:μ,:ω), Tuple{AbstractVector, AbstractMatrix}}, x)
+@inline function logdensity(d::Affine{(:μ,:ω), Tuple{AbstractVector, AbstractMatrix}}, x)
     z = x - d.μ
     lmul!(d.ω, z)
     logdensity(d.parent, z)
@@ -103,9 +167,17 @@ end
 
 logjac(d::Affine) = logjac(getfield(d, :f))
 
-
-function Base.rand(rng::Random.AbstractRNG, ::Type{T}, d::Affine) where {T}
-    z = rand(rng, T, parent(d))
+function Random.rand!(rng::Random.AbstractRNG, d::Affine, x::AbstractVector{T}, z=Vector{T}(undef, size(getfield(d,:f),2))) where {T}
+    rand!(rng, parent(d), z)
     f = getfield(d, :f)
-    return f(z)
+    apply!(x, f, z)
+    return x
 end
+
+
+# function Base.rand(rng::Random.AbstractRNG, ::Type{T}, d::Affine) where {T}
+#     f = getfield(d, :f)
+#     z = rand(rng, T, parent(d))
+#     apply!(x, f, z)
+#     return z
+# end

src/combinators/factoredbase.jl

@@ -7,7 +7,7 @@ struct FactoredBase{R,C,V,B} <: AbstractMeasure
     base :: B
 end
 
-function logdensity(d::FactoredBase, x)
+@inline function logdensity(d::FactoredBase, x)
     d.inbounds(x) || return -Inf
     d.constℓ + d.varℓ()
 end

src/combinators/pointwise.jl

@@ -27,7 +27,7 @@ Base.length(m::PointwiseProductMeasure{T}) where {T} = length(m.data)
 
 ⊙(args...) = pointwiseproduct(args...)
 
-function logdensity(d::PointwiseProductMeasure, x)
+@inline function logdensity(d::PointwiseProductMeasure, x)
     sum((logdensity(dⱼ, x) for dⱼ in d.data))
 end
 

src/combinators/power.jl

@@ -78,7 +78,7 @@ params(::Type{P}) where {F,P<:ProductMeasure{F,<:Fill}} = params(D)
 end
 
 # Same as PowerMeasure
-function logdensity(d::ProductMeasure{F,<:Fill}, x) where {F}
+@inline function logdensity(d::ProductMeasure{F,<:Fill}, x) where {F}
     d1 = d.f(first(d.pars))
     sum(xj -> logdensity(d1, xj), x)
 end

src/combinators/product.jl

@@ -62,7 +62,7 @@ function Base.show(io::IO, μ::ProductMeasure{F,T}) where {F,T <: Tuple}
     print(io, join(string.(marginals(μ)), " ⊗ "))
 end
 
-function logdensity(d::ProductMeasure{F,T}, x::Tuple) where {F,T<:Tuple}
+@inline function logdensity(d::ProductMeasure{F,T}, x::Tuple) where {F,T<:Tuple}
     mapreduce(logdensity, +, d.f.(d.pars), x)
 end
 
@@ -192,3 +192,9 @@ end
 # function Accessors.set(d::ProductMeasure{F,T}, ::typeof(params), p::Tuple) where {F, T<:Tuple}
 #     set.(marginals(d), params, p)
 # end
+
+# function logdensity(μ::ProductMeasure, ν::ProductMeasure, x)
+#     sum(zip(marginals(μ), marginals(ν), x)) do μ_ν_x
+#         logdensity(μ_ν_x...)
+#     end
+# end

src/combinators/restricted.jl

@@ -3,12 +3,14 @@ struct RestrictedMeasure{F,M} <: AbstractMeasure
     base::M
 end
 
-function logdensity(d::RestrictedMeasure, x)
+@inline function logdensity(d::RestrictedMeasure, x)
     d.f(x) || return -Inf
+    return 0.0
 end
 
 function density(d::RestrictedMeasure, x)
     d.f(x) || return 0.0
+    return 1.0
 end
 
 basemeasure(μ::RestrictedMeasure) = μ.base

src/density.jl

@@ -98,12 +98,12 @@ Define a new measure in terms of a density `f` over some measure `base`.
 
 # TODO: `density` and `logdensity` functions for `DensityMeasure`
 
-function logdensity(μ::T, ν::T, x) where {T<:AbstractMeasure}
+@inline function logdensity(μ::T, ν::T, x) where {T<:AbstractMeasure}
     μ==ν && return 0.0
     invoke(logdensity, Tuple{AbstractMeasure, AbstractMeasure, typeof(x)}, μ, ν, x)
 end
 
-function logdensity(μ::AbstractMeasure, ν::AbstractMeasure, x)
+@inline function logdensity(μ::AbstractMeasure, ν::AbstractMeasure, x)
     α = basemeasure(μ)
     β = basemeasure(ν)
 
@@ -135,6 +135,17 @@ function logdensity(μ::AbstractMeasure, ν::AbstractMeasure, x)
     return ℓ
 end
 
+function logpdf(d::AbstractMeasure, x)
+    _logpdf(d, basemeasure(d), x, zero(Float64))
+end
+
+@inline function _logpdf(d::AbstractMeasure, β::AbstractMeasure, x, ℓ::Float64)
+    d === β && return ℓ
+    Δℓ = logdensity(d, x)
+    # @show Δℓ, d
+    _logpdf(β, basemeasure(β), x, ℓ + Δℓ)
+end
+
 logdensity(::Lebesgue, ::Lebesgue, x) = 0.0
 
 # logdensity(::Lebesgue{ℝ}, ::Lebesgue{ℝ}, x) = zero(x)

src/rand.jl

@@ -6,9 +6,19 @@ Base.rand(T::Type, μ::AbstractMeasure) = rand(Random.GLOBAL_RNG, T, μ)
 
 Base.rand(rng::AbstractRNG, d::AbstractMeasure) = rand(rng, Float64, d)
 
-@inline Random.rand!(d::AbstractMeasure, arr::AbstractArray) = rand!(GLOBAL_RNG, d, arr)
-
-
+@inline Random.rand!(d::AbstractMeasure, args...) = rand!(GLOBAL_RNG, d, args...)
+
+function Base.rand(rng::Random.AbstractRNG, ::Type{T}, d::Affine) where {T}
+    z = rand(rng, T, parent(d))
+    f = getfield(d, :f)
+    return f(z)
+end
+
+# TODO: Make this work
+# function Base.rand(rng::AbstractRNG, ::Type{T}, d::AbstractMeasure) where {T}
+#     x = testvalue(d)
+#     rand!(d, x)
+# end
 
 # struct ArraySlot{A,I}
 #     arr::A