Provide more interface convenience functions (#47)

This fleshes out the interface of GMMs as vector- or set-like
(implementing `push!` and `pop!`), and MultiGMMs as dictionary-like
(implementing `valtype`, `get`, `get!`, `delete!`, and `empty!`). This
makes it easier to write code without needing to reach into internals.
The new interface functions are used to simplify some of the other code.

Author: timholy

Project.toml

@@ -1,7 +1,7 @@
 name = "GaussianMixtureAlignment"
 uuid = "f2431ed1-b9c2-4fdb-af1b-a74d6c93b3b3"
 authors = ["Tom McGrath <[email protected]> and contributors"]
-version = "0.2.1"
+version = "0.2.2"
 
 [deps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"

src/GaussianMixtureAlignment.jl

@@ -9,7 +9,7 @@ isotropic (spherical) Gaussian distributions.
 REPL help
 =========
 
-? followed by an algorith or constructor name will print help to the terminal. See: \n
+? followed by an algorithm or constructor name will print help to the terminal. See: \n
     \t?IsotropicGaussian \n
     \t?IsotropicGMM \n
     \t?IsotropicMultiGMM \n

src/draw.jl

@@ -91,7 +91,7 @@ function plot!(gd::GMMDisplay{<:NTuple{<:Any,<:AbstractIsotropicGMM}})
     label = gd[:label][]
     for (i,gmm) in enumerate(gmms)
         col = isnothing(color) ? palette[(i-1) % len + 1] : color
-        gaussiandisplay!(gd, gmm.gaussians...; display=disp, color=col, label)
+        gaussiandisplay!(gd, gmm...; display=disp, color=col, label)
     end
     return gd
 end
@@ -103,13 +103,13 @@ function plot!(gd::GMMDisplay{<:NTuple{<:Any,<:AbstractIsotropicMultiGMM{N,T,K}}
     palette = gd[:palette][]
     allkeys = Set{K}()
     for mgmm in mgmms
-        allkeys = allkeys ∪ keys(mgmm.gmms)
+        allkeys = allkeys ∪ keys(mgmm)
     end
     len = length(allkeys)
     for (i,k) in enumerate(allkeys)
         col = isnothing(color) ? palette[(i-1) % len + 1] : color
         for mgmm in mgmms
-            haskey(mgmm.gmms, k) && gmmdisplay!(gd, mgmm.gmms[k]; display=disp, color=col, palette=palette, label=string(k))
+            haskey(mgmm, k) && gmmdisplay!(gd, mgmm[k]; display=disp, color=col, palette=palette, label=string(k))
         end
     end
     return gd

src/gogma/gmm.jl

@@ -1,4 +1,5 @@
-import Base: eltype, length, size, getindex, iterate, convert, promote_rule, keys
+import Base: eltype, keytype, valtype, length, size, getindex, iterate, convert, promote_rule,
+             keys, values, push!, pop!, empty!, haskey, get, get!, delete!
 
 # Type structure: leaving things open for adding anisotropic Gaussians and GMMs
 
@@ -41,7 +42,10 @@ iterate(gmm::AbstractSingleGMM) = iterate(gmm.gaussians)
 iterate(gmm::AbstractSingleGMM, i) = iterate(gmm.gaussians, i)
 size(gmm::AbstractSingleGMM{N,T}) where {N,T} = (length(gmm.gaussians), N)
 size(gmm::AbstractSingleGMM{N,T}, idx::Int) where {N,T} = (length(gmm.gaussians), N)[idx]
-eltype(gmm::AbstractSingleGMM) = eltype(gmm.gaussians);
+eltype(gmm::AbstractSingleGMM) = eltype(gmm.gaussians)
+push!(gmm::AbstractSingleGMM, g::AbstractGaussian) = push!(gmm.gaussians, g)
+pop!(gmm::AbstractSingleGMM) = pop!(gmm.gaussians)
+empty!(gmm::AbstractSingleGMM) = empty!(gmm.gaussians)
 
 coords(gmm::AbstractSingleGMM) = hcat([g.μ for g in gmm.gaussians]...)
 weights(gmm::AbstractSingleGMM) = [g.ϕ for g in gmm.gaussians]
@@ -54,7 +58,16 @@ iterate(mgmm::AbstractMultiGMM) = iterate(mgmm.gmms)
 iterate(mgmm::AbstractMultiGMM, i) = iterate(mgmm.gmms, i)
 size(mgmm::AbstractMultiGMM{N,T,K}) where {N,T,K} = (length(mgmm.gmms), N)
 size(mgmm::AbstractMultiGMM{N,T,K}, idx::Int) where {N,T,K} = (length(mgmm.gmms), N)[idx]
-eltype(mgmm::AbstractMultiGMM) = eltype(mgmm.gmms);
+eltype(mgmm::AbstractMultiGMM) = eltype(mgmm.gmms)
+eltype(::Type{MGMM}) where MGMM<:AbstractMultiGMM = Pair{keytype(MGMM),valtype(MGMM)}
+keytype(mgmm::AbstractMultiGMM) = keytype(typeof(mgmm))
+keytype(::Type{<:AbstractMultiGMM{N,T,K}}) where {N,T,K} = K
+valtype(mgmm::AbstractMultiGMM) = valtype(mgmm.gmms)
+haskey(mgmm::AbstractMultiGMM, k) = haskey(mgmm.gmms, k)
+get(mgmm::AbstractMultiGMM, k, default) = get(mgmm.gmms, k, default)
+get!(::Type{V}, mgmm::AbstractMultiGMM, k) where V = get!(V, mgmm.gmms, k)
+delete!(mgmm::AbstractMultiGMM, k) = delete!(mgmm.gmms, k)
+empty!(mgmm::AbstractMultiGMM) = empty!(mgmm.gmms)
 
 coords(mgmm::AbstractMultiGMM) = hcat([coords(gmm) for (k,gmm) in mgmm.gmms]...)
 weights(mgmm::AbstractMultiGMM) = vcat([weights(gmm) for (k,gmm) in mgmm.gmms]...)
@@ -79,7 +92,7 @@ end
 
 IsotropicGaussian(g::AbstractIsotropicGaussian) = IsotropicGaussian(g.μ, g.σ, g.ϕ)
 
-convert(::Type{IsotropicGaussian{N,T}}, g::AbstractIsotropicGaussian) where {N,T} = IsotropicGaussian(SVector{N,T}(g.μ), T(g.σ), T(g.ϕ))
+convert(::Type{IsotropicGaussian{N,T}}, g::AbstractIsotropicGaussian) where {N,T} = IsotropicGaussian{N,T}(g.μ, g.σ, g.ϕ)
 promote_rule(::Type{IsotropicGaussian{N,T}}, ::Type{IsotropicGaussian{N,S}}) where {N,T<:Real,S<:Real} = IsotropicGaussian{N,promote_type(T,S)}
 
 (g::IsotropicGaussian)(pos::AbstractVector) = exp(-sum(abs2, pos-g.μ)/(2*g.σ^2))*g.ϕ
@@ -92,11 +105,13 @@ struct IsotropicGMM{N,T} <: AbstractIsotropicGMM{N,T}
 end
 
 IsotropicGMM(gmm::AbstractIsotropicGMM) = IsotropicGMM(gmm.gaussians)
+IsotropicGMM{N,T}() where {N,T} = IsotropicGMM{N,T}(IsotropicGaussian{N,T}[])
 
-convert(t::Type{IsotropicGMM}, gmm::AbstractIsotropicGMM) = t(gmm.gaussians)
+convert(::Type{GMM}, gmm::AbstractIsotropicGMM) where GMM<:IsotropicGMM = GMM(gmm.gaussians)
 promote_rule(::Type{IsotropicGMM{N,T}}, ::Type{IsotropicGMM{N,S}}) where {T,S,N} = IsotropicGMM{N,promote_type(T,S)}
+eltype(::Type{IsotropicGMM{N,T}}) where {N,T} = IsotropicGaussian{N,T}
 
-(gmm::IsotropicGMM)(pos::AbstractVector) = sum(g(pos) for g in gmm.gaussians)
+(gmm::IsotropicGMM)(pos::AbstractVector) = sum(g(pos) for g in gmm)
 
 """
 A collection of labeled `IsotropicGMM`s, to each be considered separately during an alignment procedure. That is,
@@ -110,6 +125,7 @@ IsotropicMultiGMM(gmm::AbstractIsotropicMultiGMM) = IsotropicMultiGMM(gmm.gmms)
 
 convert(t::Type{IsotropicMultiGMM}, mgmm::AbstractIsotropicMultiGMM) = t(mgmm.gmms)
 promote_rule(::Type{IsotropicMultiGMM{N,T,K}}, ::Type{IsotropicMultiGMM{N,S,K}}) where {N,T,S,K} = IsotropicMultiGMM{N,promote_type(T,S),K}
+valtype(::Type{IsotropicMultiGMM{N,T,K}}) where {N,T,K} = IsotropicGMM{N,T}
 
 # descriptive display
 # TODO update to display type parameters, make use of supertypes, etc

test/runtests.jl

@@ -96,6 +96,50 @@ end
     end
 end
 
+@testset "GMM interface" begin
+    tetrahedral = [
+        [0.,0.,1.],
+        [sqrt(8/9), 0., -1/3],
+        [-sqrt(2/9),sqrt(2/3),-1/3],
+        [-sqrt(2/9),-sqrt(2/3),-1/3]
+    ]
+    ch_g = IsotropicGaussian(tetrahedral[1], 1.0, 1.0)
+    s_gs = [IsotropicGaussian(x, 0.5, 1.0) for (i,x) in enumerate(tetrahedral)]
+    gmm = IsotropicGMM(s_gs)
+    @test length(gmm) == 4
+    @test gmm[2] == s_gs[2]
+    @test collect(gmm) == s_gs       # tests iterate
+    @test eltype(gmm) === eltype(typeof(gmm)) === IsotropicGaussian{3,Float64}
+    @test convert(IsotropicGMM{3,Float32}, gmm) isa IsotropicGMM{3,Float32}
+    @test_throws DimensionMismatch convert(IsotropicGMM{2,Float64}, gmm)
+    mgmmx = IsotropicMultiGMM(Dict(
+        :positive => IsotropicGMM([ch_g]),
+        :steric => gmm
+    ))
+    @test keys(mgmmx) == Set([:positive, :steric])
+    @test keytype(mgmmx) === keytype(typeof(mgmmx)) === Symbol
+    @test valtype(mgmmx) === valtype(typeof(mgmmx)) === IsotropicGMM{3,Float64}
+    @test eltype(mgmmx) === eltype(typeof(mgmmx)) === Pair{Symbol, IsotropicGMM{3,Float64}}
+    @test length(mgmmx) == 2
+    @test length(mgmmx[:steric]) == 4
+    @test mgmmx[:steric][2] == s_gs[2]
+    @test collect(mgmmx) == collect(mgmmx.gmms) # tests iterate
+    @test get!(valtype(mgmmx), mgmmx, :positive) == mgmmx[:positive]
+    gmm = get!(valtype(mgmmx), mgmmx, :acceptor)
+    @test isempty(gmm) && gmm isa IsotropicGMM{3,Float64}
+    push!(gmm, ch_g)
+    @test length(gmm) == 1
+    pop!(gmm)
+    @test isempty(gmm)
+    push!(gmm, ch_g)
+    empty!(gmm)
+    @test isempty(gmm)
+    delete!(mgmmx, :acceptor)
+    @test !haskey(mgmmx, :acceptor)
+    empty!(mgmmx)
+    @test isempty(mgmmx)
+end
+
 @testset "bounds for shrinking searchspace around an optimum" begin
     # two sets of points, each forming a 3-4-5 triangle
     xpts = [[0.,0.,0.], [3.,0.,0.,], [0.,4.,0.]]