Update MvNormal (#160)

* update MvNormal to use Affine

* typos

* inline _logpdf

* make things easily dubuggable

* update tests

* AND ANOTHER

* update tests, but MANY NEED FIXING

* LogExpFunctions 0.3.3 (for xlog1py)

* working on tests

* logpdf from MeasureBase

* logpdf test

* Symbolic-friendly StudentT

* drop redudant code

* basekernel

* bugfix

* Returns{T} <: Function

* fixing stuff

* asparams(::Affine)

* debugging, and tests

* use Measurebase.Returns (usually from Base) instead of making another

* adjust for Productmesaure updates

* Some refactoring

* MvNormal updates

* compat `replace`

* oops typo

* Actually, let's just make this `replace` stuff more direct

* bump version

Author: cscherrer

Project.toml

@@ -1,7 +1,7 @@
 name = "MeasureTheory"
 uuid = "eadaa1a4-d27c-401d-8699-e962e1bbc33b"
 authors = ["Chad Scherrer <[email protected]> and contributors"]
-version = "0.12.0"
+version = "0.13.0"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
@@ -37,7 +37,7 @@ DynamicIterators = "0.4"
 FillArrays = "0.12"
 InfiniteArrays = "0.11"
 KeywordCalls = "0.2"
-LogExpFunctions = "0.3"
+LogExpFunctions = "0.3.3"
 MLStyle = "0.4"
 MacroTools = "0.5"
 MappedArrays = "0.4"

src/MeasureTheory.jl

@@ -28,6 +28,7 @@ export CountingMeasure
 export TrivialMeasure
 export Likelihood
 export testvalue
+export basekernel
 
 using InfiniteArrays
 using ConcreteStructs
@@ -39,7 +40,12 @@ using StatsFuns
 using SpecialFunctions
 using LogExpFunctions 
 
+import NamedTupleTools
+
 import MeasureBase: testvalue, logdensity, density, basemeasure, kernel, params, ∫
+import MeasureBase: affine, supportdim
+
+import Base: rand
 
 using Reexport
 @reexport using MeasureBase
@@ -53,38 +59,22 @@ export as
 export Affine
 export AffineTransform
 
-if VERSION < v"1.7.0-beta1.0"
-    @eval begin
-        struct Returns{T}
-            value::T
-        end
-
-        (f::Returns)(x) = f.value
-    end
-end
+using MeasureBase: Returns
 
 sampletype(μ::AbstractMeasure) = typeof(testvalue(μ))
 
 # sampletype(μ::AbstractMeasure) = sampletype(basemeasure(μ))
 
-import Distributions: pdf, logpdf
 
 
-export pdf, logpdf
-
-function logpdf(d::AbstractMeasure, x)
-    _logpdf(d, x, logdensity(d,x))
-end
 
-function _logpdf(d::AbstractMeasure, x, acc::Real)
-    β = basemeasure(d)
-    d === β && return acc
+import Distributions: logpdf, pdf
 
-    _logpdf(β, x, acc + logdensity(β, x))
-end
+export pdf, logpdf
 
+Distributions.logpdf(d::AbstractMeasure, x) = MeasureBase.logpdf(d, x)
 
-pdf(d::AbstractMeasure, x) = exp(logpdf(d, x))
+Distributions.pdf(d::AbstractMeasure, x) = exp(Dists.logpdf(d, x))
 
 """
     logdensity(μ::AbstractMeasure [, ν::AbstractMeasure], x::X)
@@ -95,16 +85,26 @@ is understood to be `basemeasure(μ)`.
 """
 function logdensity end
 
+
+const AFFINEPARS = [
+    (:μ,:σ)
+    (:μ,:ω)
+    (:σ,)
+    (:ω,)
+    (:μ,)
+]
+
+
 include("const.jl")
 # include("traits.jl")
 include("parameterized.jl")
 # include("resettablerng.jl")
 
+include("combinators/affine.jl")
 include("combinators/weighted.jl")
 include("combinators/product.jl")
 include("combinators/transforms.jl")
 include("combinators/chain.jl")
-# include("combinators/basemeasure.jl")
 
 include("distributions.jl")
 

src/combinators/affine.jl

@@ -0,0 +1,7 @@
+asparams(::Affine, ::Val{:μ}) = asℝ
+asparams(::Affine, ::Val{:σ}) = asℝ₊
+asparams(::Type{A}, ::Val{:μ}) where {A<:Affine} = asℝ
+asparams(::Type{A}, ::Val{:σ}) where {A<:Affine} = asℝ₊
+
+asparams(::Affine{N,M,T}, ::Val{:μ}) where {N,M,T<:AbstractArray} = as(Array, asℝ, size(d.μ))
+asparams(::Affine{N,M,T}, ::Val{:σ}) where {N,M,T<:AbstractArray} = as(Array, asℝ, size(d.σ))

src/combinators/product.jl

@@ -1,20 +1,19 @@
-
-
-
-function TV.as(d::ProductMeasure{F,A}) where {F,A<:AbstractArray}
+function TV.as(d::ProductMeasure{F,S,A}) where {F,S,A<:AbstractArray}
     d1 = marginals(d).f(first(marginals(d).data))
     as(Array, as(d1), size(marginals(d))...)
 end
 
 ###############################################################################
 # I <: Base.Generator
 
-function TV.as(d::ProductMeasure{F,I}) where {F, I<:Base.Generator}
+function TV.as(d::ProductMeasure{F,S,I}) where {F,S,I<:Base.Generator}
     d1 = marginals(d).f(first(marginals(d).iter))
     as(Array, as(d1), size(marginals(d))...) 
 end
 
-
+function TV.as(d::ProductMeasure{Returns{T},F,A}) where {T, F, A <: AbstractArray}
+    as(Array, as(d.f.f.value), size(d.pars))
+end
 
 function Base.rand(rng::AbstractRNG, ::Type{T}, d::ProductMeasure, d1::Dists.Distribution) where {T}
     mar = marginals(d)
@@ -32,11 +31,11 @@ end
 
 
 # e.g. set(Normal(μ=2)^5, params, randn(5))
-function Accessors.set(d::ProductMeasure{F,A}, ::typeof(params), p::AbstractArray) where {F,A<:AbstractArray}
+function Accessors.set(d::ProductMeasure{F,S,A}, ::typeof(params), p::AbstractArray) where {F,S,A<:AbstractArray}
     set.(marginals(d), params, p)
 end
 
-function Accessors.set(d::ProductMeasure{F,A}, ::typeof(params), p) where {F,A<:AbstractArray}
+function Accessors.set(d::ProductMeasure{F,S,A}, ::typeof(params), p) where {F,S,A<:AbstractArray}
     par = typeof(d.pars[1])(p)
     ProductMeasure(d.f, Fill(par, size(d.pars)))
 end

src/distproxy.jl

@@ -26,7 +26,7 @@ for m in keys(PROXIES)
 end
 
 
-Base.rand(rng::AbstractRNG, T::Type, d::ParameterizedMeasure) = rand(rng, distproxy(d))
+Base.rand(rng::AbstractRNG, ::Type{T}, d::ParameterizedMeasure) where {T} = rand(rng, distproxy(d))
 
 # MonteCarloMeasurements.Particles(N::Int, d::AbstractMeasure) = MonteCarloMeasurements.Particles(N, distproxy(d))
 

src/parameterized.jl

@@ -74,7 +74,3 @@ asparams(μ::ParameterizedMeasure, nt::NamedTuple=NamedTuple()) = asparams(const
 
 TV.as(::Half) = asℝ₊
 
-asparams(::Affine, ::Val{:μ}) = asℝ
-asparams(::Affine, ::Val{:σ}) = asℝ₊
-asparams(::Type{A}, ::Val{:μ}) where {A<:Affine} = asℝ
-asparams(::Type{A}, ::Val{:σ}) where {A<:Affine} = asℝ₊

src/parameterized/cauchy.jl

@@ -6,14 +6,23 @@ export Cauchy, HalfCauchy
 @parameterized Cauchy(μ,σ) ≃ (1/π) * Lebesgue(ℝ)
 
 @kwstruct Cauchy()
+@kwstruct Cauchy(μ)
+@kwstruct Cauchy(σ)
+@kwstruct Cauchy(μ,σ)
+@kwstruct Cauchy(ω)
+@kwstruct Cauchy(μ,ω)
 
-Cauchy(nt::NamedTuple{(:μ,:σ)}) = Affine(nt, Cauchy())
-Cauchy(nt::NamedTuple{(:μ,:ω)}) = Affine(nt, Cauchy())
-Cauchy(nt::NamedTuple{(:σ,)}) = Affine(nt, Cauchy())
-Cauchy(nt::NamedTuple{(:ω,)}) = Affine(nt, Cauchy())
-Cauchy(nt::NamedTuple{(:μ,)}) = Affine(nt, Cauchy())
 
-@affinepars Cauchy
+
+for N in AFFINEPARS
+    @eval begin
+        proxy(d::Cauchy{$N}) = affine(params(d), Cauchy())
+        logdensity(d::Cauchy{$N}, x) = logdensity(proxy(d), x)
+        basemeasure(d::Cauchy{$N}) = basemeasure(proxy(d))
+    end
+end
+
+# @affinepars Cauchy
 
 function logdensity(d::Cauchy{()} , x) 
     return -log1p(x^2)
@@ -34,3 +43,8 @@ TV.as(::Cauchy) = asℝ
 HalfCauchy(σ) = HalfCauchy(σ=σ)
 
 distproxy(d::Cauchy{()}) = Dists.Cauchy()
+distproxy(d::Cauchy{(:μ,)}) = Dists.Cauchy(d.μ, 1.0)
+distproxy(d::Cauchy{(:σ,)}) = Dists.Cauchy(0.0, d.σ)
+distproxy(d::Cauchy{(:μ,:σ)}) = Dists.Cauchy(d.μ, d.σ)
+distproxy(d::Cauchy{(:ω,)}) = Dists.Cauchy(0.0, inv(d.ω))
+distproxy(d::Cauchy{(:μ,:ω)}) = Dists.Cauchy(d.μ, inv(d.ω))

src/parameterized/gumbel.jl

@@ -2,14 +2,25 @@
 
 export Gumbel
 
-@parameterized Gumbel() ≃ Lebesgue(ℝ)
+@parameterized Gumbel() 
+
+basemeasure(::Gumbel{()}) = Lebesgue(ℝ)
 
 @kwstruct Gumbel()
 
-Gumbel(nt::NamedTuple{(:σ,)}) = Affine(nt, Gumbel())
+@kwstruct Gumbel(β)
 
-@affinepars Gumbel
+@kwstruct Gumbel(μ,β)
 
+# map affine names to those more common for Gumbel
+for N in [(:μ,), (:σ,), (:μ,:σ)]
+    G = tuple(replace(collect(N), :σ => :β)...)
+    @eval begin
+        proxy(d::Gumbel{$G}) = affine(NamedTuple{$N}(values(params(d))), Gumbel())
+        logdensity(d::Gumbel{$G}, x) = logdensity(proxy(d), x)
+        basemeasure(d::Gumbel{$G}) = basemeasure(proxy(d))
+    end
+end
 
 function logdensity(d::Gumbel{()} , x)
     return -exp(-x) - x

src/parameterized/laplace.jl

@@ -5,13 +5,15 @@ export Laplace
 
 @parameterized Laplace() ≪ (1/2) * Lebesgue(ℝ)
 
-Laplace(nt::NamedTuple{(:μ,:σ)}) = Affine(nt, Laplace())
-Laplace(nt::NamedTuple{(:μ,:ω)}) = Affine(nt, Laplace())
-Laplace(nt::NamedTuple{(:σ,)}) = Affine(nt, Laplace())
-Laplace(nt::NamedTuple{(:ω,)}) = Affine(nt, Laplace())
-Laplace(nt::NamedTuple{(:μ,)}) = Affine(nt, Laplace())
+for N in AFFINEPARS
+    @eval begin
+        proxy(d::Laplace{$N}) = affine(params(d), Laplace())
+        logdensity(d::Laplace{$N}, x) = logdensity(proxy(d), x)
+        basemeasure(d::Laplace{$N}) = basemeasure(proxy(d))
+    end
+end
 
-@affinepars Laplace
+# @affinepars Laplace
 
 
 function logdensity(d::Laplace{()} , x)
@@ -25,3 +27,8 @@ Base.rand(rng::AbstractRNG, μ::Laplace{()}) = rand(rng, Dists.Laplace())
 TV.as(::Laplace) = asℝ
 
 distproxy(::Laplace{()}) = Dists.Laplace()
+distproxy(d::Laplace{(:μ,)}) = Dists.Laplace(d.μ, 1.0)
+distproxy(d::Laplace{(:σ,)}) = Dists.Laplace(0.0, d.σ)
+distproxy(d::Laplace{(:μ,:σ)}) = Dists.Laplace(d.μ, d.σ)
+distproxy(d::Laplace{(:ω,)}) = Dists.Laplace(0.0, inv(d.ω))
+distproxy(d::Laplace{(:μ,:ω)}) = Dists.Laplace(d.μ, inv(d.ω))

src/parameterized/mvnormal.jl

@@ -1,61 +1,47 @@
 
 # Multivariate Normal distribution
 
-using LinearAlgebra
 export MvNormal
-using Random
-import Base
 
+@parameterized MvNormal(μ,σ)
 
-struct MvNormal{N, T, I, J} <: ParameterizedMeasure{N}
-    par::NamedTuple{N, T}
-end
+# MvNormal(;kwargs...) = MvNormal(kwargs)
 
+@kwstruct MvNormal(μ)
+@kwstruct MvNormal(σ)
+@kwstruct MvNormal(ω)
+@kwstruct MvNormal(μ,σ)
+@kwstruct MvNormal(μ,ω)
 
-function MvNormal(nt::NamedTuple{N,T}) where {N,T}
-    I,J = mvnormaldims(nt)
+supportdim(d::MvNormal) = supportdim(params(d))
 
-    cache = Vector{Float64}(undef, max(I,J))
-    MvNormal{N,T,I,J}(cache, nt)
-end
+proxy(d::MvNormal) = affine(params(d), Normal() ^ supportdim(d))
+logdensity(d::MvNormal, x) = logdensity(proxy(d), x)
+basemeasure(d::MvNormal) = basemeasure(proxy(d))
 
-function Base.size(d::MvNormal{N, T, I, J}) where {N,T,I,J}
-    return (I,)
-end
+rand(rng::AbstractRNG, ::Type{T}, d::MvNormal) where {T} = rand(rng, T, proxy(d))
 
-mvnormaldims(nt::NamedTuple{(:A, :b)}) = size(nt.A)
-
-function MeasureTheory.basemeasure(μ::MvNormal{N, T, I,I}) where {N, T, I}
-    return (1 / sqrt2π)^I * Lebesgue(ℝ)^I
-end
-
-sampletype(d::MvNormal{N, T, I, J}) where {N,T,I,J} = Vector{Float64}
-
-MvNormal(; kwargs...) = begin
-    MvNormal((; kwargs...))
-end
-
-
-
-function Random.rand!(rng::AbstractRNG, d::MvNormal{(:A, :b),T,I,J}, x::AbstractArray) where {T,I,J}
-    z = getcache(d, J)
-    rand!(rng, Normal()^J, z)
-    copyto!(x, d.b)
-    mul!(x, d.A, z, 1.0, 1.0)
-    return x
-end
+# function MvNormal(nt::NamedTuple{(:μ,)})
+#     dim = size(nt.μ)
+#     affine(nt, Normal() ^ dim)
+# end
 
-function logdensity(d::MvNormal{(:A,:b)}, x)
-    cache = getcache(d, size(d))
-    z = d.A \ (x - d.b)
-    return logdensity(MvNormal(), z) - logabsdet(d.A)
-end
+# function MvNormal(nt::NamedTuple{(:σ,)})
+#     dim = colsize(nt.σ)
+#     affine(nt, Normal() ^ dim)
+# end
 
-≪(::MvNormal, ::Lebesgue{ℝ}) = true
+# function MvNormal(nt::NamedTuple{(:ω,)})
+#     dim = rowsize(nt.ω)
+#     affine(nt, Normal() ^ dim)
+# end
 
-# function logdensity(d::MvNormal{(:Σ⁻¹,)}, x)
-#     @tullio ℓ = -0.5 * x[i] * d.Σ⁻¹[i,j] * x[j]
-#     return ℓ
+# function MvNormal(nt::NamedTuple{(:μ, :σ,)})
+#     dim = colsize(nt.σ)
+#     affine(nt, Normal() ^ dim)
 # end
 
-mvnormaldims(nt::NamedTuple{(:Σ⁻¹,)}) = size(nt.Σ⁻¹)
+# function MvNormal(nt::NamedTuple{(:μ, :ω,)})
+#     dim = rowsize(nt.ω)
+#     affine(nt, Normal() ^ dim)
+# end