ForwardDiff support in solve_RLM and use precision matrices

Author: Affie

IncrementalInference/Project.toml

@@ -20,6 +20,7 @@ DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 FunctionalStateMachine = "3e9e306e-7e3c-11e9-12d2-8f8f67a2f951"
 IncrementalInferenceTypes = "9808408f-4dbc-47e4-913c-6068b950e289"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
@@ -83,6 +84,7 @@ FileIO = "1"
 FiniteDiff = "2"
 FiniteDifferences = "0.12"
 Flux = "0.14, 0.15, 0.16"
+ForwardDiff = "1.3.3"
 FunctionalStateMachine = "0.2.9, 0.3"
 Gadfly = "1"
 IncrementalInferenceTypes = "0.1.0"

IncrementalInference/src/manifolds/services/ManifoldsExtentions.jl

@@ -55,6 +55,18 @@ function Manifolds.get_vector!(M::NPowerManifold, Y, p, c, B::AbstractBasis)
   return Y
 end
 
+# Allocating get_vector that infers element type from coordinates (Dual-compatible)
+function ManifoldsBase.get_vector(M::NPowerManifold, p, c, B::AbstractBasis)
+  dim = manifold_dimension(M.manifold)
+  rep_size = Manifolds.representation_size(M.manifold)
+  v_iter = Ref(0)
+  return [begin
+    coords_i = SVector{dim}(view(c, v_iter[]+1:v_iter[]+dim))
+    v_iter[] += dim
+    get_vector(M.manifold, Manifolds._read(M, rep_size, p, i), coords_i, B)
+  end for i in Manifolds.get_iterator(M)]
+end
+
 function Manifolds.exp!(M::NPowerManifold, q, p, X)
   rep_size = Manifolds.representation_size(M.manifold)
   for i in Manifolds.get_iterator(M)
@@ -67,6 +79,16 @@ function Manifolds.exp!(M::NPowerManifold, q, p, X)
   return q
 end
 
+# Allocating exp that infers element type from tangent vector (Dual-compatible)
+function ManifoldsBase.exp(M::NPowerManifold, p, X)
+  rep_size = Manifolds.representation_size(M.manifold)
+  return [exp(
+    M.manifold,
+    Manifolds._read(M, rep_size, p, i),
+    Manifolds._read(M, rep_size, X, i),
+  ) for i in Manifolds.get_iterator(M)]
+end
+
 function LieGroups.compose!(M::NPowerManifold, x, p, q)
   rep_size = representation_size(M.manifold)
   for i in Manifolds.get_iterator(M)

IncrementalInference/src/parametric/services/ParametricManopt.jl

@@ -3,7 +3,7 @@ using FiniteDiff
 using SparseDiffTools
 using SparseArrays
 
-# using ForwardDiff
+using ForwardDiff
 # using Zygote
 
 ##
@@ -111,7 +111,7 @@ end
 function calcFactorResVec!(
   x::Vector{T},
   cfm_part::Vector{<:CalcFactorResidual{FT, N, D}},
-  p::AbstractArray{T},
+  p::AbstractArray,
   st::Int
 ) where {T, FT, N, D}
   for cfm in cfm_part
@@ -121,7 +121,7 @@ function calcFactorResVec!(
   return st
 end
 
-function calcFactorResVec_threaded!(x::Vector{T}, cfm_part::Vector{<:CalcFactorResidual}, p::AbstractArray{T}, st::Int) where T
+function calcFactorResVec_threaded!(x::Vector{T}, cfm_part::Vector{<:CalcFactorResidual}, p::AbstractArray, st::Int) where T
   l = getDimension(cfm_part[1]) # all should be the same
   N = length(cfm_part)
   chunkies = Iterators.partition(1:N, N ÷ Threads.nthreads())
@@ -135,7 +135,7 @@ function calcFactorResVec_threaded!(x::Vector{T}, cfm_part::Vector{<:CalcFactorR
   return st + l*N
 end
 
-function (costf::CostFres!{CFT})(M::AbstractManifold, x::Vector{T}, p::AbstractVector{T}) where {CFT,T}
+function (costf::CostFres!{CFT})(M::AbstractManifold, x::Vector{T}, p::AbstractVector) where {CFT,T}
   st = 1
   for cfm_part in costf.costfuns.x
       # if length(cfm_part) > Threads.nthreads() * 10
@@ -150,11 +150,11 @@ end
 ## --------------------------------------------------------------------------------------------------------------
 ## jacobian of function for Riemannian Levenberg-Marquardt
 ## --------------------------------------------------------------------------------------------------------------
-struct JacF_RLM!{CF, T, JC}
+struct JacF_RLM!{CF, TX, TQ, JC}
   costF!::CF
   X0::Vector{Float64}
-  X::T
-  q::T
+  X::TX
+  q::TQ
   res::Vector{Float64}
   Jcache::JC
 end
@@ -245,6 +245,48 @@ end
   #   ManifoldDiff.default_differential_backend()
   # )
 
+## --------------------------------------------------------------------------------------------------------------
+## ForwardDiff jacobian for Riemannian Levenberg-Marquardt
+## --------------------------------------------------------------------------------------------------------------
+
+struct JacF_RLM_ForwardDiff!{CF}
+  costF!::CF
+  X0::Vector{Float64}
+  res::Vector{Float64}
+end
+
+function JacF_RLM_ForwardDiff!(M, costF!, p, _fg=nothing;
+  all_points=p,
+  basis_domain::AbstractBasis = LieGroups.DefaultLieAlgebraOrthogonalBasis(),
+  is_sparse=false,  # unused, kept for interface compat
+)
+  res = reduce(vcat, map(f -> f(all_points), Vector(costF!.costfuns)))
+  X0 = zeros(manifold_dimension(M))
+  return JacF_RLM_ForwardDiff!(costF!, X0, res)
+end
+
+function (jacF!::JacF_RLM_ForwardDiff!)(
+  M::AbstractManifold,
+  J,
+  p;
+  basis_domain::AbstractBasis = DefaultOrthogonalBasis(),
+)
+  X0 = jacF!.X0
+  fill!(X0, 0)
+
+  nres = length(jacF!.res)
+  function costf(Xc)
+    X = get_vector(M, p, Xc, basis_domain)
+    q = exp(M, p, X)
+    _res = zeros(eltype(Xc), nres)
+    jacF!.costF!(M, _res, q)
+    return _res
+  end
+
+  ForwardDiff.jacobian!(J, costf, X0)
+  return J
+end
+
 struct FactorGradient{A <: AbstractMatrix}
   manifold::AbstractManifold
   JacF!::JacF_RLM!
@@ -278,9 +320,7 @@ function getSparsityPattern(fg, varLabels, factLabels)
   return sparse(getindex.(iter,1), getindex.(iter,2), ones(Bool, length(iter)))
 end
 
-# TODO only calculate marginal covariances
-
-function covarianceFiniteDiff(M, jacF!::JacF_RLM!, p0)
+function precisionFiniteDiff(M, jacF!::JacF_RLM!, p0)
     # Jcache
     X0 = fill!(deepcopy(jacF!.X0), 0)
 
@@ -292,17 +332,22 @@ function covarianceFiniteDiff(M, jacF!::JacF_RLM!, p0)
       end
     end
 
-    H = FiniteDiff.finite_difference_hessian(costf, X0)
-
-    # inv(H)
-    Σ = try 
-        Matrix(H) \ Matrix{eltype(H)}(I, size(H)...)
-      catch ex #TODO only catch correct exception and try with pinv as fallback in certain cases.
-        @warn "Hessian inverse failed" ex H
-        pinv(H)
-        # nothing
-      end
-    return Σ
+    FiniteDiff.finite_difference_hessian(costf, X0)
+end
+
+function precisionFiniteDiff(M, jacF!::JacF_RLM_ForwardDiff!, p0)
+    X0 = fill!(copy(jacF!.X0), 0)
+    nres = length(jacF!.res)
+    
+    function costf(Xc)
+      X = get_vector(M, p0, Xc, DefaultOrthogonalBasis())
+      q = exp(M, p0, X)
+      _res = zeros(nres)
+      jacF!.costF!(M, _res, q)
+      return 1/2*norm(_res)^2
+    end
+    
+    FiniteDiff.finite_difference_hessian(costf, X0)
 end
 
 function qr_linear_subsolver!(sk, JJ, grad_f_c)
@@ -316,6 +361,7 @@ function solve_RLM(
   faclabels = lsf(fg);
   is_sparse = true,
   finiteDiffCovariance = false,
+  jacobian_method::Symbol = :finitediff,
   solveKey::Symbol = :parametric,
   linear_subsolver! = qr_linear_subsolver!,
   kwargs...
@@ -341,13 +387,17 @@ function solve_RLM(
   costF! = CostFres!(calcfacs, collect(varlabelsAP))
 
   # jacobian of function for Riemannian Levenberg-Marquardt
-  jacF! = JacF_RLM!(M, costF!, p0, fg; is_sparse)
+  if jacobian_method == :forwarddiff
+    jacF! = JacF_RLM_ForwardDiff!(M, costF!, p0)
+  else
+    jacF! = JacF_RLM!(M, costF!, p0, fg; is_sparse)
+  end
 
   num_components = length(jacF!.res)
   initial_residual_values = zeros(num_components)
 
   # initial_jacobian_f not type stable, but function barrier so should be ok.
-  initial_jacobian_f = is_sparse ? 
+  initial_jacobian_f = (jacF! isa JacF_RLM! && is_sparse) ? 
     jacF!.Jcache.sparsity : 
     zeros(num_components, manifold_dimension(M))
 
@@ -365,23 +415,15 @@ function solve_RLM(
     kwargs...
   )
 
-  if length(initial_residual_values) < 1000 
-    if finiteDiffCovariance
-      # TODO this seems to be correct, but way to slow
-      Σ = covarianceFiniteDiff(M, jacF!, lm_r)
-    else
-      # TODO make sure J initial_jacobian_f is updated, otherwise recalc jacF!(M, J, lm_r) # lm_r === p0
-      J = initial_jacobian_f
-      H = J'J # approx
-      Σ = H \ Matrix{eltype(H)}(I, size(H)...)
-      # Σ = pinv(H)
-    end
+  if finiteDiffCovariance
+    Λ = precisionFiniteDiff(M, jacF!, lm_r)
   else
-    @warn "Not estimating a Dense Covariance $(size(initial_jacobian_f))"
-    Σ = nothing  
+    J = initial_jacobian_f
+    jacF!(M, J, lm_r) # recompute J at solution point
+    Λ = Symmetric(J'J) # approx Hessian = precision matrix
   end
 
-  return M, varlabelsAP, lm_r, Σ
+  return M, varlabelsAP, lm_r, Λ
 end
 
   # nlso = NonlinearLeastSquaresObjective(
@@ -440,6 +482,7 @@ function solve_RLM_conditional(
   separators::Vector{Symbol} = setdiff(ls(fg), frontals);
   is_sparse=false,
   finiteDiffCovariance=true,
+  jacobian_method::Symbol = :finitediff,
   solveKey::Symbol = :parametric,
   kwargs...
 )
@@ -498,13 +541,17 @@ function solve_RLM_conditional(
   costF! = CostFres_cond!(all_points, calcfacs, Vector{Symbol}(collect(all_varlabelsAP)))
 
   # jacobian of function for Riemannian Levenberg-Marquardt
-  jacF! = JacF_RLM!(M, costF!, p0, fg; all_points, is_sparse)
+  if jacobian_method == :forwarddiff
+    jacF! = JacF_RLM_ForwardDiff!(M, costF!, p0; all_points)
+  else
+    jacF! = JacF_RLM!(M, costF!, p0, fg; all_points, is_sparse)
+  end
 
   num_components = length(jacF!.res)
 
   initial_residual_values = zeros(num_components)
 
-  initial_jacobian_f = is_sparse ? 
+  initial_jacobian_f = (jacF! isa JacF_RLM! && is_sparse) ? 
     jacF!.Jcache.sparsity : 
     zeros(num_components, manifold_dimension(M))
 
@@ -521,13 +568,13 @@ function solve_RLM_conditional(
   )
 
   if finiteDiffCovariance
-    Σ = covarianceFiniteDiff(M, jacF!, lm_r)
+    Λ = precisionFiniteDiff(M, jacF!, lm_r)
   else
-    J = initial_jacobian_f
-    Σ = pinv(J'J)
+    jacF!(M, initial_jacobian_f, lm_r)
+    Λ = Symmetric(initial_jacobian_f' * initial_jacobian_f)
   end
-
-  return M, frontal_varlabelsAP, lm_r, Σ
+  
+  return M, frontal_varlabelsAP, lm_r, Λ
 end
 
 function extractMarginalsAP(M, labelsAP::ArrayPartition{Symbol}, Σ::AbstractArray{<:Real})
@@ -608,12 +655,14 @@ function autoinitParametric!(
         )
       )
     end
-    M, vartypeslist, lm_r, Σ = solve_RLM_conditional(dfg, [initme], initfrom; solveKey, kwargs...)
+    M, vartypeslist, lm_r, Λ = solve_RLM_conditional(dfg, [initme], initfrom; solveKey, kwargs...)
 
     val = lm_r[1]
     DFG.refMeans(vnd)[1] = val
 
-    !isnothing(Σ) && (DFG.refCovariances(vnd)[1] .= Σ)
+    if !isnothing(Λ)
+      DFG.refCovariances(vnd)[1] .= inv(Matrix(Λ))
+    end
 
     # updateSolverDataParametric!(vnd, val, Σ)
 
@@ -656,11 +705,11 @@ function DFG.solveGraphParametric!(
     error("TODO: not implemented")
   end
 
-  M, v, r, Σ = solve_RLM(fg, args...; is_sparse, kwargs...)
+  M, v, r, Λ = solve_RLM(fg, args...; is_sparse, kwargs...)
 
-  updateParametricSolution!(fg, M, v, r, Σ)
+  updateParametricSolution!(fg, M, v, r, Λ)
 
-  return M, v, r, Σ 
+  return M, v, r, Λ 
 end