Switch to easier format for interactions (#45)

* Switch to easier format for interactions

* Added test for interactions validation

* Revisions from Tim, v0.2

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.11"
+version = "0.2"
 
 [deps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"

src/gogma/bounds.jl

@@ -1,6 +1,17 @@
 loose_distance_bounds(x::AbstractGaussian, y::AbstractGaussian, args...) = loose_distance_bounds(x.μ, y.μ, args...)
 tight_distance_bounds(x::AbstractGaussian, y::AbstractGaussian, args...) = tight_distance_bounds(x.μ, y.μ, args...)
 
+function validate_interactions(interactions::Dict{Tuple{K,K},V}) where {K,V<:Number}
+    for (k1,k2) in keys(interactions)
+        if k1 != k2
+            if haskey(interactions, (k2,k1))
+                return false
+            end
+        end
+    end
+    return true
+end
+
 # 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, interactions=nothing)
     t = promote_type(numbertype(gmmx),numbertype(gmmy))
@@ -14,26 +25,31 @@ function pairwise_consts(gmmx::AbstractIsotropicGMM, gmmy::AbstractIsotropicGMM,
     return pσ, pϕ
 end
 
-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}
+function pairwise_consts(mgmmx::AbstractMultiGMM{N,T,K}, mgmmy::AbstractMultiGMM{N,S,K}, interactions::Union{Nothing,Dict{Tuple{K,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}}()
+        interactions = Dict{Tuple{K,K},t}()
         for key in xkeys ∩ ykeys
-            interactions[key] = Dict{K, t}(key => one(t))
+            interactions[(key,key)] = one(t)
         end
+    else
+        @assert validate_interactions(interactions) "Interactions must not include redundant key pairs (i.e. (k1,k2) and (k2,k1))"
     end
     mpσ, mpϕ = Dict{K, Dict{K, Matrix{t}}}(), Dict{K, Dict{K,Matrix{t}}}()
-    for key1 in keys(interactions)
+    ukeys = unique(Iterators.flatten(keys(interactions)))
+    for key1 in ukeys
         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 
+            for key2 in ukeys
+                keypair = (key1,key2)
+                keypair = haskey(interactions, keypair) ? keypair : (key2,key1)
+                if key2 ∈ ykeys && haskey(interactions, keypair)
                     pσ, pϕ = pairwise_consts(mgmmx.gmms[key1], mgmmy.gmms[key2])
                     push!(mpσ[key1], key2 => pσ)
-                    push!(mpϕ[key1], key2 => interactions[key1][key2] .* pϕ)
+                    push!(mpϕ[key1], key2 => interactions[keypair] .* pϕ)
                 end
             end
             if isempty(mpσ[key1])

src/gogma/overlap.jl

@@ -104,8 +104,9 @@ end
 
 function force!(f::AbstractVector, x::AbstractIsotropicGaussian, y::AbstractIsotropicGMM, pσ=nothing, pϕ=nothing; kwargs...)
     if isnothing(pσ) && isnothing(pϕ)
-        pσ = x.σ^2 .+ [gy.σ^2 for gy in y.gaussians]
-        pϕ = [ x.ϕ * gy.ϕ for gy in y.gaussians]
+        xσsq = x.σ^2
+        pσ = [xσsq + gy.σ^2 for gy in y.gaussians]
+        pϕ = [x.ϕ * gy.ϕ for gy in y.gaussians]
     end
     for (gy, s, w) in zip(y.gaussians, pσ, pϕ)
         force!(f, x, gy, s, w; kwargs...)
@@ -121,10 +122,8 @@ function force!(f::AbstractVector, x::AbstractIsotropicGMM, y::AbstractIsotropic
     end
 end
 
-function force!(f::AbstractVector, x::AbstractMultiGMM, y::AbstractMultiGMM, mpσ=nothing, mpϕ=nothing; interactions=nothing)
-    if isnothing(mpσ) && isnothing(mpϕ)
-        mpσ, mpϕ = pairwise_consts(x, y, interactions) 
-    end
+function force!(f::AbstractVector, x::AbstractMultiGMM, y::AbstractMultiGMM; interactions=nothing)
+    mpσ, mpϕ = pairwise_consts(x, y, interactions) 
     for k1 in keys(mpσ)
         for k2 in keys(mpσ[k1])
             # don't pass coef as a keyword argument, since the interaction coefficient is baked into mpϕ

test/runtests.jl

@@ -158,18 +158,20 @@ end
         :negative => IsotropicGMM([ch_g]),
         :steric => IsotropicGMM(s_gs)
     ))
+    # interaction validation
     interactions = Dict(
-        :positive => Dict(
-            :positive => -1.0,
-            :negative => 1.0,
-        ),
-        :negative => Dict(
-            :positive => 1.0,
-            :negative => -1.0,
-        ),
-        :steric => Dict(
-            :steric => -1.0,
-        ),
+        (:positive, :negative) =>  1.0,
+        (:negative, :positive) => 1.0,
+        (:positive, :positive) => -1.0,
+        (:negative, :negative) => -1.0,
+        (:steric, :steric) => -1.0,
+    )
+    @test_throws AssertionError GMA.pairwise_consts(mgmmx, mgmmy, interactions)
+    interactions = Dict(
+        (:positive, :negative) =>  1.0,
+        (:positive, :positive) => -1.0,
+        (:negative, :negative) => -1.0,
+        (:steric, :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)
@@ -208,17 +210,10 @@ end
     fliptform = AffineMap(RotationVec(π,0,0),[0,0,3]) ∘ AffineMap(RotationVec(0,0,π),[0,0,0])
     mgmmy = fliptform(mgmmy)
     interactions = Dict(
-        :positive => Dict(
-            :positive => -1.0,
-            :negative => 1.0,
-        ),
-        :negative => Dict(
-            :positive => 1.0,
-            :negative => -1.0,
-        ),
-        :steric => Dict(
-            :steric => -1.0,
-        ),
+        (:positive, :negative) =>  1.0,
+        (:positive, :positive) => -1.0,
+        (:negative, :negative) => -1.0,
+        (:steric, :steric) => -1.0,
     )
     f = zeros(3)
     force!(f, mgmmx, mgmmy; interactions=interactions)