Compute the "force" between two Gaussians (#43)

This is the gradient of `overlap` with respect to the centroid of the
"probe" Gaussian.

Author: timholy

Project.toml

@@ -27,9 +27,10 @@ Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 GaussianMixtureAlignmentMakieExt = "Makie"
 
 [compat]
+Colors = "0.12"
 CoordinateTransformations = "0.6"
 Distances = "0.10"
-Colors = "0.12"
+ForwardDiff = "0.10"
 GenericLinearAlgebra = "0.3"
 GeometryBasics = "0.4"
 Hungarian = "0.7"
@@ -45,8 +46,9 @@ StaticArrays = "1.5"
 julia = "1.10"
 
 [extras]
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "IntervalSets"]
+test = ["Test", "ForwardDiff", "IntervalSets"]

src/GaussianMixtureAlignment.jl

@@ -38,7 +38,7 @@ using Colors
 
 export AbstractGaussian, AbstractGMM
 export IsotropicGaussian, IsotropicGMM, IsotropicMultiGMM
-export overlap, gogma_align, rot_gogma_align, trl_gogma_align, tiv_gogma_align
+export overlap, force!, gogma_align, rot_gogma_align, trl_gogma_align, tiv_gogma_align
 export rocs_align
 export PointSet, MultiPointSet
 export kabsch, icp, iterative_hungarian, goicp_align, goih_align, tiv_goicp_align, tiv_goih_align

src/gogma/overlap.jl

@@ -88,4 +88,28 @@ Calculates the tanimoto distance based on Gaussian overlap between two GMMs.
 function tanimoto(x::AbstractGMM, y::AbstractGMM)
     o = overlap(x,y)
     return o / (overlap(x,x) + overlap(y,y) - o)
-end
+end
+
+## Forces
+
+function force!(f::AbstractVector, x::AbstractVector, y::AbstractVector, s::Real, w::Real)
+    Δ = y - x
+    f .+= Δ / s * overlap(sum(abs2, Δ), s, w)
+end
+
+function force!(f::AbstractVector, x::AbstractIsotropicGaussian, y::AbstractIsotropicGaussian,
+                s=x.σ^2+y.σ^2, w=x.ϕ*y.ϕ; coef=1)
+        return force!(f, x.μ, y.μ, s, coef*w)
+end
+
+function force!(f::AbstractVector, x::AbstractIsotropicGaussian, y::AbstractIsotropicGMM; kwargs...)
+    for gy in y.gaussians
+        force!(f, x, gy; kwargs...)
+    end
+end
+
+function force!(f::AbstractVector, x::AbstractIsotropicGMM, y::AbstractIsotropicGMM; kwargs...)
+    for gx in x.gaussians
+        force!(f, gx, y; kwargs...)
+    end
+end

test/runtests.jl

@@ -5,6 +5,7 @@ using LinearAlgebra
 using StaticArrays
 using Rotations
 using CoordinateTransformations
+using ForwardDiff
 
 using GaussianMixtureAlignment: UncertaintyRegion, RotationRegion, TranslationRegion
 using GaussianMixtureAlignment: tight_distance_bounds, loose_distance_bounds, gauss_l2_bounds, subranges, sqrt3, UncertaintyRegion, subregions, branchbound, rocs_align, overlap, gogma_align, tiv_gogma_align, tiv_goih_align, overlapobj
@@ -174,6 +175,21 @@ end
     res = gogma_align(randtform(mgmmx), mgmmy; interactions=interactions, maxsplits=5e3, nextblockfun=GMA.randomblock)
 end
 
+@testset "Forces" begin
+    μx = randn(SVector{3,Float64})
+    μy = randn(SVector{3,Float64})
+    σx = 1 + rand()
+    σy = 1 + rand()
+    ϕx = 1 + rand()
+    ϕy = 1 + rand()
+    x = IsotropicGaussian(μx, σx, ϕx)
+    y = IsotropicGaussian(μy, σy, ϕy)
+    f = zeros(3)
+    force!(f, x, y)
+    ovlp(μ) = overlap(IsotropicGaussian(μ, σx, ϕx), y)
+    @test f ≈ ForwardDiff.gradient(ovlp, μx)
+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.]]