Merge pull request #299 from JuliaRobotics/master

release v0.8.5-rc1
JuliaRobotics · Aug 8, 2024 · 6f1afc8 · 6f1afc8
2 parents ce8c995 + 2ab8e23
commit 6f1afc8
Show file tree

Hide file tree

Showing 17 changed files with 2,919 additions and 319 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "ApproxManifoldProducts"
 uuid = "9bbbb610-88a1-53cd-9763-118ce10c1f89"
 keywords = ["MM-iSAMv2", "SLAM", "inference", "sum-product", "belief-propagation", "nonparametric", "manifolds", "functional"]
-version = "0.8.4"
+version = "0.8.5"
 
 [deps]
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
@@ -51,8 +51,9 @@ Reexport = "0.2, 1.0"
 Requires = "0.5, 1"
 Rotations = "1"
 StaticArrays = "0.15, 1"
+Statistics = "1"
 TensorCast = "0.2, 0.3, 0.4"
-TransformUtils = "0.2.10"
+TransformUtils = "0.2.17"
 julia = "1.9"
 
 [extensions]

diff --git a/src/API.jl b/src/API.jl
@@ -30,21 +30,24 @@ plot( x=getPoints(pq)[1,:], y=getPoints(pq)[2,:], Geom.histogram2d )
 # TODO, convenient plotting (work in progress...)
 ```
 """
-function manifoldProduct( ff::AbstractVector{<:ManifoldKernelDensity},
-                          mani::M=ff[1].manifold;
-                          makeCopy::Bool=false,
-                          Niter::Integer=1,
-                          # partialDimsWorkaround=1:MB.manifold_dimension(mani),
-                          ndims::Integer=maximum([0;Ndim.(ff)]),
-                          N::Integer = maximum([0;Npts.(ff)]),
-                          u0 = getPoints(ff[1], false)[1],
-                          oldPoints::AbstractVector{P}= [identity_element(mani, u0) for i in 1:N],
-                          addEntropy::Bool=true,
-                          recordLabels::Bool=false,
-                          selectedLabels::Vector{Vector{Int}}=Vector{Vector{Int}}(),
-                          _randU=Vector{Float64}(),
-                          _randN=Vector{Float64}(),
-                          logger=ConsoleLogger()  ) where {M <: MB.AbstractManifold, P}
+function manifoldProduct(
+  ff::AbstractVector{<:ManifoldKernelDensity},
+  mani::M=ff[1].manifold;
+  makeCopy::Bool=false,
+  Niter::Integer=1,
+  # partialDimsWorkaround=1:MB.manifold_dimension(mani),
+  ndims::Integer=maximum([0;Ndim.(ff)]),
+  N::Integer = maximum([0;Npts.(ff)]),
+  u0 = getPoints(ff[1], false)[1],
+  oldPoints::AbstractVector{P}= [identity_element(mani, u0) for i in 1:N],
+  addEntropy::Bool=true,
+  recordLabels::Bool=false,
+  selectedLabels::Vector{Vector{Int}}=Vector{Vector{Int}}(),
+  _randU=Vector{Float64}(),
+  _randN=Vector{Float64}(),
+  logger=ConsoleLogger(),
+  bws = 0 == length(ff) ? [1.0;] : ones(ndims),
+) where {M <: MB.AbstractManifold, P}
   #
   # check quick exit
   if  1 == length(ff)
@@ -66,7 +69,6 @@ function manifoldProduct( ff::AbstractVector{<:ManifoldKernelDensity},
   glbs = KDE.makeEmptyGbGlb();
   glbs.recordChoosen = recordLabels
 
-  bws = 0 == length(ff) ? [1.0;] : ones(ndims)
   # MAKE SURE inplace ends up as matrix of coordinates from incoming ::Vector{P}
   oldpts = _pointsToMatrixCoords(mani, oldPoints)
   # FIXME currently assumes oldPoints are in coordinates...

diff --git a/src/ApproxManifoldProducts.jl b/src/ApproxManifoldProducts.jl
@@ -1,38 +1,44 @@
 module ApproxManifoldProducts
 
-using Reexport
-@reexport using KernelDensityEstimate
+import Base: *, isapprox, convert, show, eltype, length
+
+using Logging
+using StaticArrays
+using LinearAlgebra
+import LinearAlgebra: rotate!, det
+
+using TensorCast
+using DocStringExtensions
+using Distributions
 using Random
+import Random: rand
+
+using Statistics
+import Statistics: mean, std, cov, var # , entropy? 24Q3, # JL v1.11-rc1, Statistics v1.11.1.
 
-import TransformUtils as TUs
 import Rotations as _Rot
+using CoordinateTransformations
 
-import ManifoldsBase
-import ManifoldsBase: AbstractManifold, distance
 using RecursiveArrayTools: ArrayPartition
 export ArrayPartition
 
+import ManifoldsBase
+import ManifoldsBase: AbstractManifold, distance
 using Manifolds
 
-using DocStringExtensions
-
 using NLsolve
 import Optim
-using CoordinateTransformations
-using Requires
-using LinearAlgebra
-using TensorCast
-using StaticArrays
-using Logging
-using Statistics
-using Distributions
 
-import Random: rand
+import TransformUtils as TUs
+import TransformUtils: skew
+
+using Reexport
+@reexport using KernelDensityEstimate
+import KernelDensityEstimate: getPoints, getBW, evalAvgLogL, entropy, evaluate
 
-import Base: *, isapprox, convert, show, eltype
-import LinearAlgebra: rotate!, det
-import Statistics: mean, std, cov, var
-import KernelDensityEstimate: getPoints, getBW, evalAvgLogL, entropy
+
+# FIXME ON FIRE OBSOLETE REMOVE
+using Requires
 
 const MB = ManifoldsBase
 const CTs = CoordinateTransformations
@@ -48,16 +54,22 @@ include("ExportAPI.jl")
 # internal features not exported
 include("_BiMaps.jl")
 
-# AMP types and some legacy support 
+include("entities/KernelEval.jl")
+include("entities/ManellicTree.jl") # experimental
 include("entities/ManifoldKernelDensity.jl")
+
+include("services/ManifoldsOverloads.jl")
+include("CommonUtils.jl")
+include("services/ManellicTree.jl")
+
+# AMP types and some legacy support 
 include("entities/ManifoldDefinitions.jl")
 include("Legacy.jl")
 include("services/ManifoldPartials.jl")
 include("Interface.jl")
 
 # regular features
 include("services/KernelEval.jl")
-include("CommonUtils.jl")
 include("services/ManifoldKernelDensity.jl")
 include("services/Euclidean.jl")
 include("services/CircularUtils.jl")
@@ -71,10 +83,6 @@ include("API.jl")
 
 include("Deprecated.jl")
 
-# experimental
-include("entities/ManellicTree.jl")
-include("services/ManellicTree.jl")
-
 # weak dependencies
 include("../ext/WeakdepsPrototypes.jl")
 

diff --git a/src/CommonUtils.jl b/src/CommonUtils.jl
@@ -103,6 +103,33 @@ end
 #   return var(mkd.manifold, getPoints(mkd); kwargs...)
 # end
 
+_makevec(w::AbstractVector) = w
+_makevec(w::Tuple) = [w...]
+
+
+function calcProductGaussians_flat(
+  M::AbstractManifold, 
+  μ_::Union{<:AbstractVector{P},<:NTuple{N,P}}, # point type commonly known as P (actually on-manifold)
+  Σ_::Union{<:AbstractVector{S},<:NTuple{N,S}};
+  μ0 = mean(M, _makevec(μ_)), # Tangent space reference around the evenly weighted mean of incoming points
+  Λ_ = inv.(Σ_),
+  weight::Real = 1.0,
+  do_transport_correction::Bool = true
+) where {N,P<:AbstractArray,S<:AbstractMatrix{<:Real}}
+  # calc sum of covariances  
+  Λ = +(Λ_...)
+
+  # calc the covariance weighted delta means of incoming points and covariances
+  ΛΔμc = mapreduce(+, zip(Λ_, μ_)) do (s,u)
+    Δuvee = vee(M, μ0, log(M, μ0, u))
+    s*Δuvee
+  end
+
+  # calculate the delta mean
+  Δμc = Λ \ ΛΔμc
+
+  return Δμc, inv(Λ)
+end
 
 """
     $SIGNATURES
@@ -116,52 +143,98 @@ Notes
 - calc lambdas first and use to calculate mean product second.
 - https://ccrma.stanford.edu/~jos/sasp/Product_Two_Gaussian_PDFs.html
 - Pennec, X. Intrinsic Statistics on Riemannian Manifolds: Basic Tools for Geometric Measurements, HAL Archive, 2011, Inria, France.
+
+DevNotes:
+- FIXME is parallel transport needed as products involve covariances from different tangent spaces?
+- TODO avoid recomputing covariance matrix inverses all the time
 """
-function calcProductGaussians(M::AbstractManifold, 
-                              μ_::Union{<:AbstractVector{P},<:NTuple{N,P}}, # point type commonly known as P 
-                              Σ_::Union{Nothing,<:AbstractVector{S},<:NTuple{N,S}};
-                                dim::Integer=manifold_dimension(M),
-                              Λ_ = inv.(Σ_),
-                              ) where {N,P,S<:AbstractMatrix{<:Real}}
-  #
-  # calc sum of covariances  
-  Λ = zeros(MMatrix{dim,dim})
-  Λ .= sum(Λ_)
+function calcProductGaussians(
+  M::AbstractManifold, 
+  μ_::Union{<:AbstractVector{P},<:NTuple{N,P}}, # point type commonly known as P (actually on-manifold)
+  Σ_::Union{<:AbstractVector{S},<:NTuple{N,S}};
+  μ0 = mean(M, _makevec(μ_)), # Tangent space reference around the evenly weighted mean of incoming points
+  Λ_ = inv.(Σ_),
+  weight::Real = 1.0,
+  do_transport_correction::Bool = true
+) where {N,P<:AbstractArray,S<:AbstractMatrix{<:Real}}
+  # step 1, basic/naive Gaussian product (ignoring disjointed covariance coordinates) 
+  Δμn, Σn = calcProductGaussians_flat(M, μ_, Σ_; μ0, Λ_, weight)
+  Δμ = exp(M, μ0, hat(M, μ0, Δμn))
+
+  # for development and testing cases return without doing transport
+  do_transport_correction ? nothing : (return Δμ, Σn)
+
+  # first transport (push forward) covariances to common coordinates
+  # see [Ge, van Goor, Mahony, 2024]
+  iΔμ = inv(M, Δμ)
+  μi_ = map(u->Manifolds.compose(M,iΔμ,u), μ_)
+  μi_̂  = map(u->log(M,μ0,u), μi_)
+  # μi = map(u->vee(M,μ0,u), μi_̂ )
+  Ji = ApproxManifoldProducts.parallel_transport_curvature_2nd_lie.(Ref(M), μi_̂ )
+  iJi = inv.(Ji)
+  Σi_hat = map((J,S)->J*S*(J'), iJi, Σ_)
+
+  # Reset step to absorb extended μ+ coordinates into kernel on-manifold μ 
+  # consider using Δμ in place of μ0
+  Δμplusc, Σdiam = ApproxManifoldProducts.calcProductGaussians_flat(M, μi_, Σi_hat; μ0, weight)
+  Δμplus_̂  = hat(M, μ0, Δμplusc)
+  Δμplus = exp(M, μ0, Δμplus_̂ )
+  μ_plus = Manifolds.compose(M,Δμ,Δμplus)
+  Jμ = ApproxManifoldProducts.parallel_transport_curvature_2nd_lie(M, Δμplus_̂ )
+  Σ_plus = Jμ*Σdiam*(Jμ')
 
-  # Tangent space reference around the evenly weighted mean of incoming points
-  u0 = mean(M, μ_)
-
-  # calc the covariance weighted delta means of incoming points and covariances
-  ΛΔμ = zeros(MVector{dim})
-  for (s,u) in zip(Λ_, μ_)
-    # require vee as per Pennec, Caesar Ref [3.6]
-    Δuvee = vee(M, u0, log(M, u0, u))
-    ΛΔμ += s*Δuvee
-  end
-
-  # calculate the delta mean
-  Δμ = Λ \ ΛΔμ
-
   # return new mean and covariance
-  return exp(M, u0, hat(M, u0, Δμ)), inv(Λ) 
+  return μ_plus, Σ_plus 
 end
 
 # additional support case where covariances are passed as diagonal-only vectors 
 # still pass nothing, to avoid stack overflow.  Only Λ_ is needed further
-calcProductGaussians( M::AbstractManifold, 
-                      μ_::Union{<:AbstractVector{P},<:NTuple{N,P}},
-                      Σ_::Union{<:AbstractVector{S},<:NTuple{N,S}};
-                        dim::Integer=manifold_dimension(M),
-                      Λ_ = map(s->diagm( 1.0 ./ s), Σ_),
-                      ) where {N,P,S<:AbstractVector} = calcProductGaussians(M, μ_, nothing; dim=dim, Λ_=Λ_ )
+calcProductGaussians( 
+  M::AbstractManifold, 
+  μ_::Union{<:AbstractVector{P},<:NTuple{N,P}},
+  Σ_::Union{<:AbstractVector{S},<:NTuple{N,S}};
+  dim::Integer=manifold_dimension(M),
+  Λ_ = map(s->diagm( 1.0 ./ s), Σ_),
+  weight::Real = 1.0,
+  do_transport_correction::Bool = true
+) where {N,P,S<:AbstractVector} = calcProductGaussians(M, μ_, nothing; dim, Λ_, do_transport_correction )
 #
 
-calcProductGaussians( M::AbstractManifold, 
-                      μ_::Union{<:AbstractVector{P},<:NTuple{N,P}};
-                        dim::Integer=manifold_dimension(M),
-                      Λ_ = diagm.( (1.0 ./ μ_) ),
-                      ) where {N,P} = calcProductGaussians(M, μ_, nothing; dim=dim, Λ_=Λ_ )
-#
+
+
+"""
+    $SIGNATURES
+
+EXPERIMENTAL: On-manifold product of Gaussians.
+
+DevNotes
+- FIXME do product of concentrated Gaussians on Lie group (approximation):
+  - See Section 3.2 and 4 of [Ge, van Goor, Mahony: A Geometric Perspective on using Gaussian Distributions on Lie Groups, 2024].
+  - Also see upstream utils, https://juliamanifolds.github.io/Manifolds.jl/stable/features/distributions.html
+- FIXME is parallel transport needed when multiplying with covariances from difffent tangent spaces?
+"""
+function calcProductGaussians(
+  M::AbstractManifold,
+  kernels::Union{<:AbstractVector{K},NTuple{N,K}};
+  μ0 = nothing,
+  weight::Real = 1.0,
+  do_transport_correction::Bool = true
+) where {N,K <: MvNormalKernel}
+  # CHECK this should be on-manifold for points
+  μ_ = mean.(kernels) # This is a ArrayPartition which IS DEFINITELY ON MANIFOLD (we dispatch on mean)
+  Σ_ = cov.(kernels)  # on tangent
+
+  # parallel transport needed for covariances from different tangent spaces
+  _μ, _Σ = if isnothing(μ0)
+    calcProductGaussians(M, μ_, Σ_; do_transport_correction)
+  else
+    calcProductGaussians(M, μ_, Σ_; μ0, do_transport_correction)
+  end
+
+  return MvNormalKernel(_μ, _Σ, weight)
+end
+
+
 
 
 function _update!(dst::MN, src::MN) where {MN <: ManifoldKernelDensity}

diff --git a/src/Legacy.jl b/src/Legacy.jl
@@ -117,6 +117,11 @@ import KernelDensityEstimate: Ndim, Npts, getWeights, marginal
 import KernelDensityEstimate: getKDERange, getKDEMax, getKDEMean, getKDEfit
 import KernelDensityEstimate: sample, rand, resample, kld, minkld
 
+Npts(::ManellicTree{M,D,N}) where {M,D,N} = N
+Ndim(mt::ManellicTree) = manifold_dimension(mt.manifold)
+getBW(mker::MvNormalKernel) = cov(mker) |> collect
+getBW(mt::ManellicTree) = getBW.(mt.leaf_kernels)
+
 
 Ndim(x::ManifoldKernelDensity, w...;kw...) = Ndim(x.belief,w...;kw...)
 Npts(x::ManifoldKernelDensity, w...;kw...) = Npts(x.belief,w...;kw...)

diff --git a/src/entities/KernelEval.jl b/src/entities/KernelEval.jl
@@ -0,0 +1,15 @@
+
+
+abstract type AbstractKernel end
+
+@kwdef struct MvNormalKernel{P,T,M,iM} <: AbstractKernel
+  """ On-manifold point representing center (mean) of the MvNormal distrubtion """ 
+  μ::P
+  """ Zero-mean normal distrubtion with covariance """
+  p::MvNormal{T,M}
+  # TDB might already be covered in p.Σ.chol but having issues with SymPD (not particular to this AMP repo)
+  """ Manually maintained square root concentration matrix for faster compute, TODO likely duplicate of existing Distrubtions.jl functionality. """
+  sqrt_iΣ::iM = sqrt(inv(cov(p)))
+  """ Nonparametric weight value """
+  weight::Float64 = 1.0
+end
diff --git a/src/entities/ManellicTree.jl b/src/entities/ManellicTree.jl
@@ -25,7 +25,12 @@ struct ManellicTree{M,D<:AbstractVector,N,HL,HT}
   # kernels::ArrayPartition{<:Number,KT}
   leaf_kernels::SizedVector{N,HL}
   tree_kernels::SizedVector{N,HT}
-  left_idx::MVector{N,Int}
-  right_idx::MVector{N,Int}
+  segments::SizedVector{N,Set{Int}}
+  # left_idx::MVector{N,Int}
+  # right_idx::MVector{N,Int}
+
+  # workaround to overcome bug for StaticArrays `isdefined() != false` issue
+  _workaround_isdef_leafkernel::Set{Int}
+  _workaround_isdef_treekernel::Set{Int}
 end
 
diff --git a/src/entities/ManifoldKernelDensity.jl b/src/entities/ManifoldKernelDensity.jl
@@ -1,4 +1,5 @@
 
+const TreeDensity = Union{<:ManellicTree, <: BallTreeDensity}
 
 """
     $TYPEDEF
@@ -12,7 +13,7 @@ Notes
 DevNotes
 - WIP AMP issue 41, use generic retractions during manifold products.
 """
-struct ManifoldKernelDensity{M <: MB.AbstractManifold, B <: BallTreeDensity, L, P <: AbstractArray}
+struct ManifoldKernelDensity{M <: MB.AbstractManifold, B <: TreeDensity, L, P <: AbstractArray}
   manifold::M
   """ legacy expects matrix of coordinates (as columns) """
   belief::B