Mvnormal testvalue (#207)

* fix insupport(d::MvNormal, x)

* insupport(d::OrthoLebesgue, x)

* update tests

* bump version

* MvNormal work

* formatting

* MvNormal updates

* add tests

* inline

* type stability

* drop dead code

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.16.1"
+version = "0.16.2"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

src/combinators/affine.jl

@@ -41,6 +41,11 @@ Base.size(f::AffineTransform{(:μ, :λ)}) = size(f.λ)
 Base.size(f::AffineTransform{(:σ,)}) = size(f.σ)
 Base.size(f::AffineTransform{(:λ,)}) = size(f.λ)
 
+LinearAlgebra.rank(f::AffineTransform{(:σ,)})    = rank(f.σ)
+LinearAlgebra.rank(f::AffineTransform{(:λ,)})    = rank(f.λ)
+LinearAlgebra.rank(f::AffineTransform{(:μ,:σ,)}) = rank(f.σ)
+LinearAlgebra.rank(f::AffineTransform{(:μ,:λ,)}) = rank(f.λ)
+
 function Base.size(f::AffineTransform{(:μ,)})
     (n,) = size(f.μ)
     return (n, n)
@@ -132,9 +137,16 @@ struct OrthoLebesgue{N,T} <: PrimitiveMeasure
     OrthoLebesgue(nt::NamedTuple{N,T}) where {N,T} = new{N,T}(nt)
 end
 
+function insupport(d::OrthoLebesgue, x)
+    f = AffineTransform(d.par)
+    finv = inverse(f)
+    z = finv(x)
+    f(z) ≈ x
+end
+
 basemeasure(d::OrthoLebesgue) = d
 
-logdensity_def(::OrthoLebesgue, x) = static(0)
+logdensity_def(::OrthoLebesgue, x) = static(0.0)
 
 struct Affine{N,M,T} <: AbstractMeasure
     f::AffineTransform{N,T}
@@ -147,7 +159,7 @@ function Pretty.tile(d::Affine)
     Pretty.list_layout([pars, Pretty.tile(d.parent)]; prefix = :Affine)
 end
 
-function testvalue(d::Affine)
+@inline function testvalue(d::Affine)
     f = getfield(d, :f)
     z = testvalue(parent(d))
     return f(z)

src/parameterized/mvnormal.jl

@@ -8,42 +8,77 @@ export MvNormal
 # MvNormal(;kwargs...) = MvNormal(kwargs)
 
 @kwstruct MvNormal(μ)
+
 @kwstruct MvNormal(σ)
 @kwstruct MvNormal(λ)
+
 @kwstruct MvNormal(μ, σ)
 @kwstruct MvNormal(μ, λ)
 
-supportdim(d::MvNormal) = supportdim(params(d))
+@kwstruct MvNormal(Σ)
+@kwstruct MvNormal(Λ)
+@kwstruct MvNormal(μ, Σ)
+@kwstruct MvNormal(μ, Λ)
 
-@useproxy MvNormal
+as(d::MvNormal{(:μ,)}) = as(Array, length(d.μ))
 
-proxy(d::MvNormal) = affine(params(d), Normal()^supportdim(d))
+as(d::MvNormal{(:Σ,),Tuple{C}}) where {C<:Cholesky}    = as(Array, size(d.Σ, 1))
+as(d::MvNormal{(:Λ,),Tuple{C}}) where {C<:Cholesky}    = as(Array, size(d.Λ, 1))
+as(d::MvNormal{(:μ, :Σ),Tuple{C}}) where {C<:Cholesky} = as(Array, size(d.Σ, 1))
+as(d::MvNormal{(:μ, :Λ),Tuple{C}}) where {C<:Cholesky} = as(Array, size(d.Λ, 1))
 
-rand(rng::AbstractRNG, ::Type{T}, d::MvNormal) where {T} = rand(rng, T, proxy(d))
+function as(d::MvNormal{(:σ,),Tuple{M}}) where {M<:Triangular}
+    σ = d.σ
+    if @inbounds all(i -> σ[i] > 0, diagind(σ))
+        return as(Array, size(σ, 1))
+    else
+        @error "Not implemented yet"
+    end
+end
+
+function as(d::MvNormal{(:λ,),Tuple{M}}) where {M<:Triangular}
+    λ = d.λ
+    if @inbounds all(i -> λ[i] > 0, diagind(λ))
+        return as(Array, size(λ, 1))
+    else
+        @error "Not implemented yet"
+    end
+end
 
-insupport(::MvNormal, x) = true
+for N in setdiff(AFFINEPARS, [(:μ,)])
+    @eval begin
+        function as(d::MvNormal{$N})
+            p = proxy(d)
+            if rank(getfield(p,:f)) == only(supportdim(d))
+                return as(Array, supportdim(d))
+            else
+                @error "Not yet implemented"
+            end
+        end
+    end
+end
 
-# function MvNormal(nt::NamedTuple{(:μ,)})
-#     dim = size(nt.μ)
-#     affine(nt, Normal() ^ dim)
-# end
+supportdim(d::MvNormal) = supportdim(params(d))
+
+supportdim(nt::NamedTuple{(:Σ,)}) = size(nt.Σ, 1)
+supportdim(nt::NamedTuple{(:μ,:Σ)}) = size(nt.Σ, 1)
+supportdim(nt::NamedTuple{(:Λ,)}) = size(nt.Λ, 1)
+supportdim(nt::NamedTuple{(:μ,:Λ)}) = size(nt.Λ, 1)
+
+@useproxy MvNormal
 
-# function MvNormal(nt::NamedTuple{(:σ,)})
-#     dim = colsize(nt.σ)
-#     affine(nt, Normal() ^ dim)
-# end
+for N in [(:Σ,), (:μ,:Σ), (:Λ,), (:μ,:Λ)]
+    @eval basemeasure_depth(d::MvNormal{$N}) = static(2)
+end
 
-# function MvNormal(nt::NamedTuple{(:λ,)})
-#     dim = rowsize(nt.λ)
-#     affine(nt, Normal() ^ dim)
-# end
+proxy(d::MvNormal) = affine(params(d), Normal()^supportdim(d))
+
+rand(rng::AbstractRNG, ::Type{T}, d::MvNormal) where {T} = rand(rng, T, proxy(d))
 
-# function MvNormal(nt::NamedTuple{(:μ, :σ,)})
-#     dim = colsize(nt.σ)
-#     affine(nt, Normal() ^ dim)
-# end
+insupport(d::MvNormal, x) = insupport(proxy(d), x)
 
-# function MvNormal(nt::NamedTuple{(:μ, :λ,)})
-#     dim = rowsize(nt.λ)
-#     affine(nt, Normal() ^ dim)
-# end
+# Note: (C::Cholesky).L may or may not make a copy, depending on C.uplo, which is not included in the type
+@inline proxy(d::MvNormal{(:Σ,),Tuple{C}}) where {C<:Cholesky}    = affine((σ = d.Σ.L,), Normal()^supportdim(d))
+@inline proxy(d::MvNormal{(:Λ,),Tuple{C}}) where {C<:Cholesky}    = affine((λ = d.Λ.L,), Normal()^supportdim(d))
+@inline proxy(d::MvNormal{(:μ, :Σ),Tuple{C}}) where {C<:Cholesky} = affine((μ = d.μ, σ = d.Σ.L), Normal()^supportdim(d))
+@inline proxy(d::MvNormal{(:μ, :Λ),Tuple{C}}) where {C<:Cholesky} = affine((μ = d.μ, λ = d.Λ.L), Normal()^supportdim(d))

src/utils.jl

@@ -104,3 +104,19 @@ function func_string(f, types)
         return string(f)
     end
 end
+
+function getL(C::Cholesky)
+    Cfactors = getfield(C, :factors)
+    Cuplo    = getfield(C, :uplo)
+
+    LowerTriangular(Cuplo === 'L' ? Cfactors : Cfactors')
+end
+
+function getU(C::Cholesky)
+    Cfactors = getfield(C, :factors)
+    Cuplo    = getfield(C, :uplo)
+
+    UpperTriangular(Cuplo === 'U' ? Cfactors : Cfactors')
+end
+
+const Triangular = Union{L,U} where {L<:LowerTriangular,U<:UpperTriangular}

test/runtests.jl

@@ -23,9 +23,11 @@ function draw2(μ)
     return (x, y)
 end
 
-function test_measure(μ)
-    logdensity_def(μ, testvalue(μ)) isa AbstractFloat
-end
+x = randn(10,3)
+Σ = cholesky(x'*x)
+Λ = cholesky(inv(Σ))
+σ = MeasureTheory.getL(Σ)
+λ = MeasureTheory.getL(Λ)
 
 test_measures = Any[
     # Chain(x -> Normal(μ=x), Normal(μ=0.0))
@@ -52,6 +54,12 @@ test_measures = Any[
     Normal(2, 3)
     Poisson(3.1)
     StudentT(ν = 2.1)
+    MvNormal(σ = [1 0; 0 1; 1 1])
+    MvNormal(λ = [1 0 1; 0 1 1])
+    MvNormal(Σ = Σ)
+    MvNormal(Λ = Λ)
+    MvNormal(σ = σ)
+    MvNormal(λ = λ)
     Uniform()
     Counting(Float64)
     Dirac(0.0) + Normal()
@@ -67,7 +75,6 @@ testbroken_measures = Any[
 @testset "testvalue" begin
     for μ in test_measures
         @info "testing $μ"
-        @test test_measure(μ)
         test_interface(μ)
     end