Added support for repulsion and interaction maps for MultiGMMs (#39)

* Added support for repulsion and interaction maps for MultiGMMs

* v0.1.9

Author: tmcgrath325

Project.toml

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

src/distancebounds.jl

@@ -7,29 +7,29 @@ function infbounds(x,y)
     return (typeinf, typeinf) 
 end
 
-function loose_distance_bounds(x::SVector{3,<:Number}, y::SVector{3,<:Number}, σᵣ::Number, σₜ::Number)
+function loose_distance_bounds(x::SVector{3,<:Number}, y::SVector{3,<:Number}, σᵣ::Number, σₜ::Number, maximize::Bool = false)
     ubdist = norm(x - y)
     γₜ = sqrt3 * σₜ 
     γᵣ = 2 * sin(min(sqrt3 * σᵣ, π) / 2) * norm(x)
-    lb, ub = max(ubdist - γₜ - γᵣ, 0), ubdist
+    lb, ub = maximize ? (max(ubdist - γₜ - γᵣ, 0), ubdist) : (ubddist + γₜ + γᵣ, ubdist)
     numtype = promote_type(typeof(lb), typeof(ub))
     return numtype(lb), numtype(ub)
 end
-loose_distance_bounds(x::SVector{3}, y::SVector{3}, R::RotationVec, T::SVector{3}, σᵣ, σₜ
-    ) = (R.sx^2 + R.sy^2 + R.sz^2) > pisq ? infbounds(x,y) : loose_distance_bounds(R*x, y-T, σᵣ, σₜ) # loose_distance_bounds(R*x, y-T, σᵣ, σₜ)
-loose_distance_bounds(x::SVector{3}, y::SVector{3}, block::UncertaintyRegion) = loose_distance_bounds(x, y, block.R, block.T, block.σᵣ, block.σₜ)
-loose_distance_bounds(x::SVector{3}, y::SVector{3}, block::SearchRegion) = loose_distance_bounds(x, y, UncertaintyRegion(block))
+loose_distance_bounds(x::SVector{3}, y::SVector{3}, R::RotationVec, T::SVector{3}, σᵣ, σₜ, maximize::Bool = false,
+    ) = (R.sx^2 + R.sy^2 + R.sz^2) > pisq ? infbounds(x,y) : loose_distance_bounds(R*x, y-T, σᵣ, σₜ, maximize) # loose_distance_bounds(R*x, y-T, σᵣ, σₜ)
+loose_distance_bounds(x::SVector{3}, y::SVector{3}, block::UncertaintyRegion, maximize::Bool = false) = loose_distance_bounds(x, y, block.R, block.T, block.σᵣ, block.σₜ, maximize)
+loose_distance_bounds(x::SVector{3}, y::SVector{3}, block::SearchRegion, maximize::Bool = false) = loose_distance_bounds(x, y, UncertaintyRegion(block), maximize)
 
 
 """
     lb, ub = tight_distance_bounds(x::SVector{3,<:Number}, y::SVector{3,<:Number}, σᵣ::Number, σₜ::Number)
     lb, ub = tight_distance_bounds(x::SVector{3,<:Number}, y::SVector{3,<:Number}, R::RotationVec, T<:SVector{3}, σᵣ::Number, σₜ::Number)
 
-Within an uncertainty region, find the bounds on distance between two points x and y.
+Within an uncertainty region, find the bounds on distance between two points x and y. 
 
 See [Campbell & Peterson, 2016](https://arxiv.org/abs/1603.00150)
 """
-function tight_distance_bounds(x::SVector{3,<:Number}, y::SVector{3,<:Number}, σᵣ::Number, σₜ::Number)
+function tight_distance_bounds(x::SVector{3,<:Number}, y::SVector{3,<:Number}, σᵣ::Number, σₜ::Number, maximize::Bool = false)
     # prepare positions and angles
     xnorm, ynorm = norm(x), norm(y)
     if xnorm*ynorm == 0
@@ -42,13 +42,23 @@ function tight_distance_bounds(x::SVector{3,<:Number}, y::SVector{3,<:Number}, 
     # upper bound distance at hypercube center
     ubdist = norm(x - y)
 
-    # lower bound distance from the nearest point on the "spherical cap"
-    if cosα >= cosβ
-        lbdist = max(abs(xnorm-ynorm) - sqrt3*σₜ, 0)
+    if maximize
+        # this case is intended for situations where the objective function scales negatively with distance\
+        # lbdist, which will be the further point on the spherical cap, will be larger than ubdist
+        if cosα + cosβ >= π
+            lbdist = xnorm + ynorm + sqrt3*σₜ
+        else
+            lbdist = √(xnorm^2 + ynorm^2 - 2*xnorm*ynorm*(cosα*cosβ-√((1-cosα^2)*(1-cosβ^2)))) + sqrt3*σₜ
+        end
     else
-        lbdist = try max(√(xnorm^2 + ynorm^2 - 2*xnorm*ynorm*(cosα*cosβ+√((1-cosα^2)*(1-cosβ^2)))) - sqrt3*σₜ, 0)  # law of cosines
-        catch e     # when the argument for the square root is negative (within machine precision of 0, usually)
-            0
+        # lower bound distance from the nearest point on the "spherical cap"
+        if cosα >= cosβ
+            lbdist = max(abs(xnorm-ynorm) - sqrt3*σₜ, 0)
+        else
+            lbdist = try max(√(xnorm^2 + ynorm^2 - 2*xnorm*ynorm*(cosα*cosβ+√((1-cosα^2)*(1-cosβ^2)))) - sqrt3*σₜ, 0)  # law of cosines
+            catch e     # when the argument for the square root is negative (within machine precision of 0, usually)
+                0
+            end
         end
     end
 
@@ -57,7 +67,7 @@ function tight_distance_bounds(x::SVector{3,<:Number}, y::SVector{3,<:Number}, 
     return (numtype(lbdist), numtype(ubdist))
 end
 
-tight_distance_bounds(x::SVector{3,<:Number}, y::SVector{3,<:Number}, R::RotationVec, T::SVector{3}, σᵣ::Number, σₜ::Number
-    ) = (R.sx^2 + R.sy^2 + R.sz^2) > pisq ? infbounds(x,y) : tight_distance_bounds(R*x, y-T, σᵣ, σₜ) # tight_distance_bounds(R*x, y-T, σᵣ, σₜ) 
-tight_distance_bounds(x::SVector{3}, y::SVector{3}, block::UncertaintyRegion) = tight_distance_bounds(x, y, block.R, block.T, block.σᵣ, block.σₜ)
-tight_distance_bounds(x::SVector{3}, y::SVector{3}, block::Union{RotationRegion, TranslationRegion}) = tight_distance_bounds(x, y, UncertaintyRegion(block))
+tight_distance_bounds(x::SVector{3,<:Number}, y::SVector{3,<:Number}, R::RotationVec, T::SVector{3}, σᵣ::Number, σₜ::Number, maximize::Bool = false,
+    ) = (R.sx^2 + R.sy^2 + R.sz^2) > pisq ? infbounds(x,y) : tight_distance_bounds(R*x, y-T, σᵣ, σₜ, maximize) # tight_distance_bounds(R*x, y-T, σᵣ, σₜ) 
+tight_distance_bounds(x::SVector{3}, y::SVector{3}, block::UncertaintyRegion, maximize::Bool = false) = tight_distance_bounds(x, y, block.R, block.T, block.σᵣ, block.σₜ, maximize)
+tight_distance_bounds(x::SVector{3}, y::SVector{3}, block::Union{RotationRegion, TranslationRegion}, maximize::Bool = false) = tight_distance_bounds(x, y, UncertaintyRegion(block), maximize)

src/draw.jl

@@ -21,7 +21,7 @@ const DEFAULT_COLORS = [            # CUD colors: https://jfly.uni-koeln.de/colo
 const cosθs = [cos(θ) for θ in θs]
 const sinθs = [sin(θ) for θ in θs]
 
-equal_volume_radius(σ, ϕ) = (EQUAL_VOL_CONST*ϕ)^(1/3) * σ
+equal_volume_radius(σ, ϕ) = (EQUAL_VOL_CONST*abs(ϕ))^(1/3) * σ
 
 function flat_circle!(f, pos, r, dim::Int; kwargs...)
     if dim == 3
@@ -67,7 +67,7 @@ function plot!(gd::GaussianDisplay{<:NTuple{<:Any, <:AbstractIsotropicGaussian}}
     color = gd[:color][]
     plotfun = disp == :wire ? wire_sphere! : ( disp == :solid ? solid_sphere! : throw(ArgumentError("Unrecognized display option: `$disp`")))
     for g in gauss
-        plotfun(gd, g.μ, equal_volume_radius(g.σ, g.ϕ); color=color)
+        plotfun(gd, g.μ, g.σ; color=color)
     end
     return gd
 end

src/gogma/align.jl

@@ -1,17 +1,17 @@
-function gogma_align(gmmx::AbstractGMM, gmmy::AbstractGMM; kwargs...)
-    pσ, pϕ = pairwise_consts(gmmx,gmmy)
+function gogma_align(gmmx::AbstractGMM, gmmy::AbstractGMM; interactions=nothing, kwargs...)
+    pσ, pϕ = pairwise_consts(gmmx,gmmy,interactions)
     boundsfun(x,y,block) = gauss_l2_bounds(x,y,block,pσ,pϕ)
     localfun(x,y,block) = local_align(x,y,block,pσ,pϕ)
     return branchbound(gmmx, gmmy; boundsfun=boundsfun, localfun=localfun, kwargs...)
 end
-function rot_gogma_align(gmmx::AbstractGMM, gmmy::AbstractGMM; kwargs...)
-    pσ, pϕ = pairwise_consts(gmmx,gmmy)
+function rot_gogma_align(gmmx::AbstractGMM, gmmy::AbstractGMM; interactions=nothing, kwargs...)
+    pσ, pϕ = pairwise_consts(gmmx,gmmy,interactions)
     boundsfun(x,y,block) = gauss_l2_bounds(x,y,block,pσ,pϕ)
     localfun(x,y,block) = local_align(x,y,block,pσ,pϕ)
     rot_branchbound(gmmx, gmmy; boundsfun=boundsfun, localfun=localfun, kwargs...)
 end
-function trl_gogma_align(gmmx::AbstractGMM, gmmy::AbstractGMM; kwargs...) 
-    pσ, pϕ = pairwise_consts(gmmx,gmmy)
+function trl_gogma_align(gmmx::AbstractGMM, gmmy::AbstractGMM; interactions=nothing, kwargs...) 
+    pσ, pϕ = pairwise_consts(gmmx,gmmy,interactions)
     boundsfun(x,y,block) = gauss_l2_bounds(x,y,block,pσ,pϕ)
     localfun(x,y,block) = local_align(x,y,block,pσ,pϕ)
     trl_branchbound(gmmx, gmmy; boundsfun=boundsfun, localfun=localfun, kwargs...)

src/gogma/bounds.jl

@@ -2,7 +2,7 @@ loose_distance_bounds(x::AbstractGaussian, y::AbstractGaussian, args...) = loose
 tight_distance_bounds(x::AbstractGaussian, y::AbstractGaussian, args...) = tight_distance_bounds(x.μ, y.μ, args...)
 
 # prepare pairwise values for `σx^2 + σy^2` and `ϕx * ϕy` for all gaussians in `gmmx` and `gmmy`
-function pairwise_consts(gmmx::AbstractIsotropicGMM, gmmy::AbstractIsotropicGMM)
+function pairwise_consts(gmmx::AbstractIsotropicGMM, gmmy::AbstractIsotropicGMM, interactions=nothing)
     t = promote_type(numbertype(gmmx),numbertype(gmmy))
     pσ, pϕ = zeros(t, length(gmmx), length(gmmy)), zeros(t, length(gmmx), length(gmmy))
     for (i,gaussx) in enumerate(gmmx.gaussians)
@@ -14,13 +14,33 @@ function pairwise_consts(gmmx::AbstractIsotropicGMM, gmmy::AbstractIsotropicGMM)
     return pσ, pϕ
 end
 
-function pairwise_consts(mgmmx::AbstractMultiGMM{N,T,K}, mgmmy::AbstractMultiGMM{N,S,K}) where {N,T,S,K}
-    t = promote_type(numbertype(mgmmx),numbertype(mgmmy))
-    mpσ, mpϕ = Dict{K, Matrix{t}}(), Dict{K, Matrix{t}}()
-    for key in keys(mgmmx.gmms) ∩ keys(mgmmy.gmms)
-        pσ, pϕ = pairwise_consts(mgmmx.gmms[key], mgmmy.gmms[key])
-        push!(mpσ, Pair(key, pσ))
-        push!(mpϕ, Pair(key, pϕ))
+function pairwise_consts(mgmmx::AbstractMultiGMM{N,T,K}, mgmmy::AbstractMultiGMM{N,S,K}, interactions::Union{Nothing,Dict{K,Dict{K,V}}}=nothing) where {N,T,S,K,V <: Number}
+    t = promote_type(numbertype(mgmmx),numbertype(mgmmy), isnothing(interactions) ? numbertype(mgmmx) : V)
+    xkeys = keys(mgmmx.gmms)
+    ykeys = keys(mgmmy.gmms)
+    if isnothing(interactions)
+        interactions = Dict{K, Dict{K, t}}()
+        for key in xkeys ∩ ykeys
+            interactions[key] = Dict{K, t}(key => one(t))
+        end
+    end
+    mpσ, mpϕ = Dict{K, Dict{K, Matrix{t}}}(), Dict{K, Dict{K,Matrix{t}}}()
+    for key1 in keys(interactions)
+        if key1 ∈ xkeys 
+            push!(mpσ, key1 => Dict{K, Matrix{t}}())
+            push!(mpϕ, key1 => Dict{K, Matrix{t}}())
+            for key2 in keys(interactions[key1])
+                if key2 ∈ ykeys 
+                    pσ, pϕ = pairwise_consts(mgmmx.gmms[key1], mgmmy.gmms[key2])
+                    push!(mpσ[key1], key2 => pσ)
+                    push!(mpϕ[key1], key2 => interactions[key1][key2] .* pϕ)
+                end
+            end
+            if isempty(mpσ[key1])
+                delete!(mpσ, key1)
+                delete!(mpϕ, key1)
+            end
+        end
     end
     return mpσ, mpϕ
 end
@@ -41,7 +61,7 @@ the uncertainty region is assumed to be centered at the origin (i.e. x has alrea
 See [Campbell & Peterson, 2016](https://arxiv.org/abs/1603.00150)
 """
 function gauss_l2_bounds(x::AbstractIsotropicGaussian, y::AbstractIsotropicGaussian, R::RotationVec, T::SVector{3}, σᵣ, σₜ, s=x.σ^2 + y.σ^2, w=x.ϕ*y.ϕ; distance_bound_fun = tight_distance_bounds)
-    (lbdist, ubdist) = distance_bound_fun(R*x.μ, y.μ-T, σᵣ, σₜ)
+    (lbdist, ubdist) = distance_bound_fun(R*x.μ, y.μ-T, σᵣ, σₜ, w < 0)
 
     # evaluate objective function at each distance to get upper and lower bounds
     return -overlap(lbdist^2, s, w), -overlap(ubdist^2, s, w)
@@ -58,7 +78,7 @@ gauss_l2_bounds(x::AbstractGaussian, y::AbstractGaussian, block::SearchRegion, s
 
 
 
-function gauss_l2_bounds(gmmx::AbstractSingleGMM, gmmy::AbstractSingleGMM, R::RotationVec, T::SVector{3}, σᵣ::Number, σₜ::Number, pσ=nothing, pϕ=nothing; kwargs...)
+function gauss_l2_bounds(gmmx::AbstractSingleGMM, gmmy::AbstractSingleGMM, R::RotationVec, T::SVector{3}, σᵣ::Number, σₜ::Number, pσ=nothing, pϕ=nothing, interactions=nothing; kwargs...)
     # prepare pairwise widths and weights, if not provided
     if isnothing(pσ) || isnothing(pϕ)
         pσ, pϕ = pairwise_consts(gmmx, gmmy)
@@ -75,17 +95,19 @@ function gauss_l2_bounds(gmmx::AbstractSingleGMM, gmmy::AbstractSingleGMM, R::Ro
     return lb, ub
 end
 
-function gauss_l2_bounds(mgmmx::AbstractMultiGMM, mgmmy::AbstractMultiGMM, R::RotationVec, T::SVector{3}, σᵣ::Number, σₜ::Number, mpσ=nothing, mpϕ=nothing)
+function gauss_l2_bounds(mgmmx::AbstractMultiGMM, mgmmy::AbstractMultiGMM, R::RotationVec, T::SVector{3}, σᵣ::Number, σₜ::Number, mpσ=nothing, mpϕ=nothing, interactions=nothing)
     # prepare pairwise widths and weights, if not provided
     if isnothing(mpσ) || isnothing(mpϕ)
-        mpσ, mpϕ = pairwise_consts(mgmmx, mgmmy)
+        mpσ, mpϕ = pairwise_consts(mgmmx, mgmmy, interactions)
     end
 
     # sum bounds for each pair of points
     lb = 0.
     ub = 0.
-    for key in keys(mgmmx.gmms) ∩ keys(mgmmy.gmms)
-        lb, ub = (lb, ub) .+ gauss_l2_bounds(mgmmx.gmms[key], mgmmy.gmms[key], R, T, σᵣ, σₜ, mpσ[key], mpϕ[key])
+    for (key1, intrs) in mpσ
+        for (key2, pσ) in intrs
+            lb, ub = (lb, ub) .+ gauss_l2_bounds(mgmmx.gmms[key1], mgmmy.gmms[key2], R, T, σᵣ, σₜ, pσ, mpϕ[key1][key2])
+        end
     end
     return lb, ub
 end

src/gogma/overlap.jl

@@ -61,16 +61,18 @@ end
 
 Calculates the unnormalized overlap between two `AbstractMultiGMM` objects.
 """
-function overlap(x::AbstractMultiGMM, y::AbstractMultiGMM, mpσ=nothing, mpϕ=nothing)
+function overlap(x::AbstractMultiGMM, y::AbstractMultiGMM, mpσ=nothing, mpϕ=nothing, interactions=nothing)
     # prepare pairwise widths and weights, if not provided
     if isnothing(mpσ) || isnothing(mpϕ)
-        mpσ, mpϕ = pairwise_consts(x, y)
+        mpσ, mpϕ = pairwise_consts(x, y, interactions)
     end
 
     # sum overlaps from each keyed pairs of GMM
     ovlp = zero(promote_type(numbertype(x),numbertype(y)))
-    for k in keys(x.gmms) ∩ keys(y.gmms)
-        ovlp += overlap(x.gmms[k], y.gmms[k], mpσ[k], mpϕ[k])
+    for k1 in keys(mpσ)
+        for k2 in keys(mpσ[k1])
+            ovlp += overlap(x.gmms[k1], y.gmms[k2], mpσ[k1][k2], mpϕ[k1][k2])
+        end
     end
     return ovlp
 end

src/localalign.jl

@@ -1,4 +1,21 @@
 tformwithparams(X,x) = RotationVec(X[1:3]...)*x + SVector{3}(X[4:6]...)
+# function tformwithparams(X,x) 
+#     if sum(abs2, X[1:3]) == 0 # handled for autodiff around 0
+#         T = eltype(X)
+#         θ = norm(X[1:3])
+#         a = θ > 0 ? X[1] / θ : one(T)
+#         b = θ > 0 ? X[2] / θ : zero(T)
+#         c = θ > 0 ? X[3] / θ : zero(T)
+#         R = AngleAxis(θ, a, b, c)
+#         @show R
+#     else
+#         R = RotationVec(X[1:3]...)
+#     end
+#     t = SVector{3}(X[4:6]...)
+#     @show (R*x)[1]
+#     return R*x + t
+# end
+
 overlapobj(X,x,y,args...) = -overlap(tformwithparams(X,x), y, args...)
 
 function distanceobj(X, x, y; correspondence = hungarian_assignment)
@@ -34,6 +51,10 @@ function local_align(x::AbstractModel, y::AbstractModel, block::SearchRegion, ar
 
     # set initial guess at the center of the block
     initial_X = center(block)
+    # if (typeof(block) <: UncertaintyRegion && sum(abs2, initial_X[1:Int(end/2)]) == 0) || (typeof(block) <: RotationRegion && sum(abs2, initial_X) == 0)
+    #     T = eltype(initial_X)
+    #     initial_X = initial_X .+ [eps(T), zeros(T, length(initial_X)-1)...]
+    # end
 
     # local optimization within the block
     f(X) = alignment_objective(X, x, y, block, args...; kwargs...)

test/runtests.jl

@@ -140,6 +140,40 @@ end
     @test isapprox(rocs_align(gmmx, gmmy).minimum, -overlap(gmmx,gmmx); atol=1E-12)
 end
 
+@testset "MultiGMMs with interactions" 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)]
+    mgmmx = IsotropicMultiGMM(Dict(
+        :positive => IsotropicGMM([ch_g]),
+        :steric => IsotropicGMM(s_gs)
+    ))
+    mgmmy = IsotropicMultiGMM(Dict(
+        :negative => IsotropicGMM([ch_g]),
+        :steric => IsotropicGMM(s_gs)
+    ))
+    interactions = Dict(
+        :positive => Dict(
+            :positive => -1.0,
+            :negative => 1.0,
+        ),
+        :negative => Dict(
+            :positive => 1.0,
+            :negative => -1.0,
+        ),
+        :steric => Dict(
+            :steric => -1.0,
+        ),
+    )
+    randtform = AffineMap(RotationVec(π*0.1rand(3)...), SVector{3}(0.1*rand(3)...))
+    res = gogma_align(randtform(mgmmx), mgmmy; interactions=interactions, maxsplits=5e3, nextblockfun=GMA.randomblock)
+end
+
 @testset "GO-ICP and GO-IH run without errors" begin
     xpts = [[0.,0.,0.], [3.,0.,0.,], [0.,4.,0.]] 
     ypts = [[1.,1.,1.], [1.,-2.,1.], [1.,1.,-3.]]