From 9d0e3e163abebb77d1ff03c46804b75754fed2c3 Mon Sep 17 00:00:00 2001
From: Bradley Solliday <bradley.solliday@skydio.com>
Date: Tue, 6 Jun 2023 11:17:45 +0000
Subject: [PATCH] [SymForce] Allow dense linearization in optimizer

Previously, we assumed all linearizations in the optimizer were sparse
linearizations. This commit just templates the linearization on the
matrix type so it can be used with a dense solver as well.

topic: dense_linearization_in_optimizer_too
relative: dense_optimization_stats
GitOrigin-RevId: 904cd470eb1dc11be2427f7098be12fd93bd3836
---
 symforce/opt/optimizer.h   | 15 ++++++++-------
 symforce/opt/optimizer.tcc | 16 ++++++++--------
 2 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/symforce/opt/optimizer.h b/symforce/opt/optimizer.h
index 0d2874644..d408798af 100644
--- a/symforce/opt/optimizer.h
+++ b/symforce/opt/optimizer.h
@@ -75,7 +75,8 @@ class Optimizer {
   using Scalar = ScalarType;
   using NonlinearSolver = NonlinearSolverType;
   using FailureReason = typename NonlinearSolver::FailureReason;
-  using Stats = OptimizationStats<typename NonlinearSolver::MatrixType>;
+  using MatrixType = typename NonlinearSolver::MatrixType;
+  using Stats = OptimizationStats<MatrixType>;
 
   /**
    * Base constructor
@@ -173,7 +174,7 @@ class Optimizer {
   /**
    * Linearize the problem around the given values
    */
-  SparseLinearization<Scalar> Linearize(const Values<Scalar>& values);
+  Linearization<MatrixType> Linearize(const Values<Scalar>& values);
 
   /**
    * Get covariances for each optimized key at the given linearization
@@ -184,10 +185,10 @@ class Optimizer {
    *
    * May not be called before either Optimize() or Linearize() has been called.
    */
-  void ComputeAllCovariances(const SparseLinearization<Scalar>& linearization,
+  void ComputeAllCovariances(const Linearization<MatrixType>& linearization,
                              std::unordered_map<Key, MatrixX<Scalar>>& covariances_by_key);
   [[deprecated("Pass covariances_by_key by reference instead")]] void ComputeAllCovariances(
-      const SparseLinearization<Scalar>& linearization,
+      const Linearization<MatrixType>& linearization,
       std::unordered_map<Key, MatrixX<Scalar>>* covariances_by_key);
 
   /**
@@ -207,11 +208,11 @@ class Optimizer {
    * exactly the set of keys requested.  `covariances_by_key` must not contain any keys that are not
    * in `keys`.
    */
-  void ComputeCovariances(const SparseLinearization<Scalar>& linearization,
+  void ComputeCovariances(const Linearization<MatrixType>& linearization,
                           const std::vector<Key>& keys,
                           std::unordered_map<Key, MatrixX<Scalar>>& covariances_by_key);
   [[deprecated("Pass covariances_by_key by reference instead")]] void ComputeCovariances(
-      const SparseLinearization<Scalar>& linearization, const std::vector<Key>& keys,
+      const Linearization<MatrixType>& linearization, const std::vector<Key>& keys,
       std::unordered_map<Key, MatrixX<Scalar>>* covariances_by_key);
 
   /**
@@ -288,7 +289,7 @@ class Optimizer {
    */
   struct ComputeCovariancesStorage {
     sym::MatrixX<Scalar> covariance;
-    Eigen::SparseMatrix<Scalar> H_damped;
+    MatrixType H_damped;
   };
 
   mutable ComputeCovariancesStorage compute_covariances_storage_;
diff --git a/symforce/opt/optimizer.tcc b/symforce/opt/optimizer.tcc
index f49494ce0..4cf1abe63 100644
--- a/symforce/opt/optimizer.tcc
+++ b/symforce/opt/optimizer.tcc
@@ -139,18 +139,18 @@ void Optimizer<ScalarType, NonlinearSolverType>::Optimize(Values<Scalar>* values
 }
 
 template <typename ScalarType, typename NonlinearSolverType>
-SparseLinearization<ScalarType> Optimizer<ScalarType, NonlinearSolverType>::Linearize(
-    const Values<Scalar>& values) {
+Linearization<typename NonlinearSolverType::MatrixType>
+Optimizer<ScalarType, NonlinearSolverType>::Linearize(const Values<Scalar>& values) {
   Initialize(values);
 
-  SparseLinearization<Scalar> linearization;
+  Linearization<MatrixType> linearization;
   linearize_func_(values, linearization);
   return linearization;
 }
 
 template <typename ScalarType, typename NonlinearSolverType>
 void Optimizer<ScalarType, NonlinearSolverType>::ComputeAllCovariances(
-    const SparseLinearization<Scalar>& linearization,
+    const Linearization<MatrixType>& linearization,
     std::unordered_map<Key, MatrixX<Scalar>>& covariances_by_key) {
   SYM_ASSERT(IsInitialized());
   nonlinear_solver_.ComputeCovariance(linearization.hessian_lower,
@@ -214,7 +214,7 @@ size_t ComputeBlockDimension(const Linearizer<Scalar>& linearizer, const std::ve
 
 template <typename ScalarType, typename NonlinearSolverType>
 void Optimizer<ScalarType, NonlinearSolverType>::ComputeAllCovariances(
-    const SparseLinearization<Scalar>& linearization,
+    const Linearization<MatrixType>& linearization,
     std::unordered_map<Key, MatrixX<Scalar>>* covariances_by_key) {
   SYM_ASSERT(covariances_by_key != nullptr);
   ComputeAllCovariances(linearization, *covariances_by_key);
@@ -222,7 +222,7 @@ void Optimizer<ScalarType, NonlinearSolverType>::ComputeAllCovariances(
 
 template <typename ScalarType, typename NonlinearSolverType>
 void Optimizer<ScalarType, NonlinearSolverType>::ComputeCovariances(
-    const SparseLinearization<Scalar>& linearization, const std::vector<Key>& keys,
+    const Linearization<MatrixType>& linearization, const std::vector<Key>& keys,
     std::unordered_map<Key, MatrixX<Scalar>>& covariances_by_key) {
   const bool same_order = internal::CheckKeyOrderMatchesLinearizerKeysStart(linearizer_, keys);
   SYM_ASSERT(same_order);
@@ -240,7 +240,7 @@ void Optimizer<ScalarType, NonlinearSolverType>::ComputeCovariances(
 
 template <typename ScalarType, typename NonlinearSolverType>
 void Optimizer<ScalarType, NonlinearSolverType>::ComputeCovariances(
-    const SparseLinearization<Scalar>& linearization, const std::vector<Key>& keys,
+    const Linearization<MatrixType>& linearization, const std::vector<Key>& keys,
     std::unordered_map<Key, MatrixX<Scalar>>* covariances_by_key) {
   SYM_ASSERT(covariances_by_key != nullptr);
   ComputeCovariances(linearization, keys, *covariances_by_key);
@@ -339,7 +339,7 @@ template <typename ScalarType, typename NonlinearSolverType>
 typename NonlinearSolverType::LinearizeFunc
 Optimizer<ScalarType, NonlinearSolverType>::BuildLinearizeFunc(const bool check_derivatives) {
   return [this, check_derivatives](const Values<Scalar>& values,
-                                   SparseLinearization<Scalar>& linearization) {
+                                   Linearization<MatrixType>& linearization) {
     linearizer_.Relinearize(values, linearization);
 
     if (check_derivatives) {