From dd78b100fe9f60fad443f59d494b2300ec805363 Mon Sep 17 00:00:00 2001
From: Johannes Terblanche <Affie@users.noreply.github.com>
Date: Thu, 19 Oct 2023 10:54:50 +0200
Subject: [PATCH] Experiment with factor gradients in solve

---
 .../parametric/services/ParametricManopt.jl   |  7 ++++-
 .../src/services/FactorGradients.jl           | 29 +++++++++++++++++++
 .../src/services/NumericalCalculations.jl     | 18 ++++++++++--
 3 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/IncrementalInference/src/parametric/services/ParametricManopt.jl b/IncrementalInference/src/parametric/services/ParametricManopt.jl
index 48678aaf..418db963 100644
--- a/IncrementalInference/src/parametric/services/ParametricManopt.jl
+++ b/IncrementalInference/src/parametric/services/ParametricManopt.jl
@@ -72,12 +72,17 @@ function CalcFactorResidualAP(
   return ArrayPartition{CalcFactorResidual, typeof(parts_tuple)}(parts_tuple)
 end
 
-function (cfm::CalcFactorResidual)(p)
+function (cfm::CalcFactorResidual)(p::Vector)
   meas = cfm.meas
   points = map(idx->p[idx], cfm.varOrderIdxs)
   return cfm.sqrt_iΣ * cfm(meas, points...)
 end
 
+function  (cfm::CalcFactorResidual)(p::ArrayPartition)
+  points = map(idx->p.x[idx], cfm.varOrderIdxs)
+  return cfm.sqrt_iΣ * cfm(cfm.meas, points...)
+end
+
 # cost function f: M->ℝᵈ for Riemannian Levenberg-Marquardt 
 struct CostFres_cond!{PT, CFT}
   points::PT
diff --git a/IncrementalInference/src/services/FactorGradients.jl b/IncrementalInference/src/services/FactorGradients.jl
index f8e67f54..8ff74d13 100644
--- a/IncrementalInference/src/services/FactorGradients.jl
+++ b/IncrementalInference/src/services/FactorGradients.jl
@@ -34,7 +34,36 @@ function factorJacobian(
   return ManifoldDiff.jacobian(M_dom, M_codom, costf, p0, backend)
 end
 
+function factorGradient(
+  cf::CalcFactorResidual,
+  M,
+  p,
+  backend = ManifoldDiff.TangentDiffBackend(ManifoldDiff.FiniteDiffBackend()),
+)
+  ManifoldDiff.gradient(M, (x) -> 1//2 * norm(cf(x))^2, p, backend)
+end
 
+function factorJacobian(
+  cf::CalcFactorResidual,
+  M_dom,
+  p,
+  backend = ManifoldDiff.TangentDiffBackend(ManifoldDiff.FiniteDiffBackend()),
+)
+  # M_dom = ProductManifold(getManifold.(fg, varlabels)...)
+  M_codom = Euclidean(manifold_dimension(getManifold(cf)))
+    
+  return ManifoldDiff.jacobian(M_dom, M_codom, cf, p, backend)
+end
+
+#
+function factorGradient(
+  cf::CalcFactorNormSq,
+  M,
+  p,
+  backend = ManifoldDiff.TangentDiffBackend(ManifoldDiff.FiniteDiffBackend()),
+)
+  ManifoldDiff.gradient(M, cf, p, backend)
+end
 
 export getCoordSizes
 export checkGradientsToleranceMask, calcPerturbationFromVariable
diff --git a/IncrementalInference/src/services/NumericalCalculations.jl b/IncrementalInference/src/services/NumericalCalculations.jl
index d77c9d86..7369e61c 100644
--- a/IncrementalInference/src/services/NumericalCalculations.jl
+++ b/IncrementalInference/src/services/NumericalCalculations.jl
@@ -71,15 +71,16 @@ function _solveLambdaNumeric(
   return r.minimizer
 end
 
-# struct OptimCalcConv end
+struct OptimCalcConv end
 # CalcFactorNormSq cost function for an input in coordinates as used by Optim.jl
-function (hypoCalcFactor::CalcFactorNormSq)(M::AbstractManifold, Xc::AbstractVector)
+function (hypoCalcFactor::CalcFactorNormSq)(::Type{OptimCalcConv}, M::AbstractManifold, Xc::AbstractVector)
   # hypoCalcFactor.manifold is the factor's manifold, not the variable's manifold that is needed here
   ϵ = getPointIdentity(M)
   X = get_vector(M, ϵ, SVector(Xc), DefaultOrthogonalBasis())
   p = exp(M, ϵ, X)
   return hypoCalcFactor(CalcConv, p)
 end
+(hypoCalcFactor::CalcFactorNormSq)(M::AbstractManifold, p) = hypoCalcFactor(OptimCalcConv, M, p)
 
 struct ManoptCalcConv end
 
@@ -117,10 +118,19 @@ function _solveLambdaNumeric(
       retraction_method = ExponentialRetraction()
     )
     return r
+  elseif false
+    r = gradient_descent(
+      M,
+      (M,x)->hypoCalcFactor(x),
+      (M, x)-> factorGradient(hypoCalcFactor, M, x),
+      u0;
+      stepsize=ConstantStepsize(0.1), 
+    )
+    return r
   end
 
   r = Optim.optimize(
-    x->hypoCalcFactor(M, x),
+    x->hypoCalcFactor(OptimCalcConv, M, x),
     X0c,
     alg
   )
@@ -394,6 +404,8 @@ function (cf::CalcFactorNormSq)(::Type{CalcConv}, x)
   res = isnothing(cf.slack) ? res : res .- cf.slack
   return sum(x->x^2, res)
 end
+#default to conv
+(cf::CalcFactorNormSq)(x) = cf(CalcConv, x)
 
 function _buildHypoCalcFactor(ccwl::CommonConvWrapper, smpid::Integer, _slack=nothing)
   # build a view to the decision variable memory