[SymForce] Reduce # of public Linearizer methods

This is in support of adding a dense optimizer implementation to
symforce.

Since the Linearizer class is responsible for knowing how to transfer
the outputs of the Factor linearization functions into the problem
linearization, a process which depends intextricably on the
representation of the linearization (whether its sparse or dense), a
dense optimizer would require a different linearizer than the
`Linearizer` class already present.

Thus, the simpler the public api of the linearizer classes, the easier
it is to swap out one for another.

This PR removes 2 methods from the `Linearizer` class:
- `CheckKeysAreContiguousAtStart`, whose implementation depends only on the
`Keys()` and `StateIndex()` (which are already public), which it moves
to the Optimizer because that is the only place it is used.
- `SplitCovariancesByKey`, which it moves to the covariance_utils.h

**Other notes**
I thought the implementation of `CheckKeysAreContiguousAtStart could be
rewritten to be easier to understand, so I did so in the first commit
(after this one).

In the second commit, I actually moved it out of `Linearizer` class and
split it into two functions: `CheckKeyOrderMatchesLinearizerKeysStart`
and `ComputeBlockDimension`.

In the third commit, I just moved `SplitCovariancesByKey` out of the
`Linearizer` class with only light modification (on account of the fact
that it was being made a function from a method).

Topic: reduce_surface_area_of_linearizer
GitOrigin-RevId: 34cc5b3a2ba3b88eb041c3430efbc84f255a3f9a

symforce/opt/internal/covariance_utils.h

@@ -5,9 +5,12 @@
 
 #pragma once
 
+#include <unordered_map>
+
 #include <Eigen/Sparse>
 
 #include <sym/util/typedefs.h>
+#include <symforce/opt/linearizer.h>
 
 #include "../sparse_schur_solver.h"
 
@@ -50,5 +53,32 @@ void ComputeCovarianceBlockWithSchurComplement(Eigen::SparseMatrix<Scalar>& hess
   schur_solver.SInvInPlace(covariance_block);
 }
 
+/**
+ * Extract covariances for optimized variables individually from the full problem covariance.  For
+ * each variable in `keys`, the returned matrix is the corresponding block from the diagonal of
+ * the full covariance matrix.  Requires that the Linearizer has already been initialized
+ */
+template <typename Scalar, typename MatrixType>
+void SplitCovariancesByKey(const Linearizer<Scalar>& linearizer, const MatrixType& covariance_block,
+                           const std::vector<Key>& keys,
+                           std::unordered_map<Key, MatrixX<Scalar>>& covariances_by_key) {
+  SYM_ASSERT(linearizer.IsInitialized());
+
+  // Fill out the covariance blocks
+  const auto& state_index = linearizer.StateIndex();
+  for (const auto& key : keys) {
+    const auto& entry = state_index.at(key.GetLcmType());
+
+    covariances_by_key[key] =
+        covariance_block.block(entry.offset, entry.offset, entry.tangent_dim, entry.tangent_dim);
+  }
+
+  // Make sure we have the expected number of keys
+  // If this fails, it might mean that you passed a covariances_by_key that contained keys from a
+  // different Linearizer or Optimizer, or previously called with a different subset of the problem
+  // keys
+  SYM_ASSERT(covariances_by_key.size() == keys.size());
+}
+
 }  // namespace internal
 }  // namespace sym

symforce/opt/linearizer.cc

@@ -95,36 +95,6 @@ void Linearizer<ScalarType>::Relinearize(const Values<Scalar>& values,
   }
 }
 
-template <typename ScalarType>
-bool Linearizer<ScalarType>::CheckKeysAreContiguousAtStart(const std::vector<Key>& keys,
-                                                           size_t* const block_dim) const {
-  SYM_ASSERT(!keys.empty());
-
-  auto full_problem_keys_iter = keys_.begin();
-  auto keys_iter = keys.begin();
-  for (; keys_iter != keys.end(); ++full_problem_keys_iter, ++keys_iter) {
-    if (full_problem_keys_iter == keys_.end()) {
-      throw std::runtime_error("Keys has extra entries that are not in the full problem");
-    }
-
-    if (*full_problem_keys_iter != *keys_iter) {
-      if (state_index_.find(keys_iter->GetLcmType()) == state_index_.end()) {
-        throw std::runtime_error("Tried to check key which is not in the full problem");
-      } else {
-        // The next key is in the problem, it's just out of order; so we return false
-        return false;
-      }
-    }
-  }
-
-  if (block_dim != nullptr) {
-    const auto& index_entry = state_index_.at(keys.back().GetLcmType());
-    *block_dim = index_entry.offset + index_entry.tangent_dim;
-  }
-
-  return true;
-}
-
 template <typename ScalarType>
 bool Linearizer<ScalarType>::IsInitialized() const {
   return initialized_;

symforce/opt/linearizer.h

@@ -54,27 +54,6 @@ class Linearizer {
    */
   void Relinearize(const Values<Scalar>& values, Linearization<Scalar>& linearization);
 
-  /**
-   * Check whether the keys in `keys` correspond 1-1 (and in the same order) with the start of the
-   * key ordering in the problem linearization
-   *
-   * If block_dim is provided, it will be filled out with the (tangent) dimension of the problem
-   * hessian and rhs which is occupied by the given keys
-   *
-   * TODO(aaron): Maybe kill this once we have general marginalization
-   */
-  bool CheckKeysAreContiguousAtStart(const std::vector<Key>& keys,
-                                     size_t* block_dim = nullptr) const;
-
-  /**
-   * Extract covariances for optimized variables individually from the full problem covariance.  For
-   * each variable in `keys`, the returned matrix is the corresponding block from the diagonal of
-   * the full covariance matrix.  Requires that the Linearizer has already been initialized
-   */
-  template <typename MatrixType>
-  void SplitCovariancesByKey(const MatrixType& covariance_block, const std::vector<Key>& keys,
-                             std::unordered_map<Key, MatrixX<Scalar>>& covariances_by_key) const;
-
   /**
    * Whether this contains values, versus having not been evaluated yet
    */
@@ -87,6 +66,9 @@ class Linearizer {
 
   const std::vector<Key>& Keys() const;
 
+  // NOTE(brad): Offset of index entries is sum of all tangent_dims of all previous index entries
+  // (order of index entries determined by order of corresponding keys in Keys()). Contains entry
+  // for each key in Keys().
   const std::unordered_map<key_t, index_entry_t>& StateIndex() const;
 
  private:
@@ -175,8 +157,6 @@ Linearization<Scalar> Linearize(const std::vector<Factor<Scalar>>& factors,
 
 }  // namespace sym
 
-#include "./linearizer.tcc"
-
 // Explicit instantiation declaration
 extern template class sym::Linearizer<double>;
 extern template class sym::Linearizer<float>;

symforce/opt/optimizer.tcc

@@ -155,10 +155,63 @@ void Optimizer<ScalarType, NonlinearSolverType>::ComputeAllCovariances(
   SYM_ASSERT(IsInitialized());
   nonlinear_solver_.ComputeCovariance(linearization.hessian_lower,
                                       compute_covariances_storage_.covariance);
-  linearizer_.SplitCovariancesByKey(compute_covariances_storage_.covariance, keys_,
-                                    covariances_by_key);
+  internal::SplitCovariancesByKey(linearizer_, compute_covariances_storage_.covariance, keys_,
+                                  covariances_by_key);
 }
 
+namespace internal {
+
+/**
+ * Check whether the keys in `keys` correspond 1-1 (and in the same order) with the start of the
+ * key ordering in the problem linearization
+ *
+ * Throws runtime_error if keys is longer than linearizer.Keys() or sometimes if key not in Keys
+ * found at start of linearizer.Keys().
+ *
+ * TODO(aaron): Maybe kill this once we have general marginalization
+ */
+template <typename Scalar>
+bool CheckKeyOrderMatchesLinearizerKeysStart(const Linearizer<Scalar>& linearizer,
+                                             const std::vector<Key>& keys) {
+  SYM_ASSERT(!keys.empty());
+
+  const std::vector<Key>& full_problem_keys = linearizer.Keys();
+  if (full_problem_keys.size() < keys.size()) {
+    throw std::runtime_error("Keys has extra entries that are not in the full problem");
+  }
+
+  const std::unordered_map<key_t, index_entry_t>& state_index = linearizer.StateIndex();
+  for (int i = 0; i < static_cast<int>(keys.size()); i++) {
+    if (keys[i] != full_problem_keys[i]) {
+      if (state_index.find(keys[i].GetLcmType()) == state_index.end()) {
+        throw std::runtime_error("Tried to check key which is not in the full problem");
+      } else {
+        // The next key is in the problem, it's just out of order; so we return false
+        return false;
+      }
+    }
+  }
+
+  return true;
+}
+
+/**
+ * Returns the (tangent) dimension of the problem hessian and rhs which is occupied by
+ * the given keys.
+ *
+ * Precondition: keys equals first keys.size() entries of linearizer.Keys()
+ */
+template <typename Scalar>
+size_t ComputeBlockDimension(const Linearizer<Scalar>& linearizer, const std::vector<Key>& keys) {
+  // The idea is that the offset of a state index entry is the sum of the tangent dims of all of
+  // the previous keys, so we just add the tangent_dim of the last key to it's offset to get the
+  // sum of all tangent dims.
+  const auto& last_index_entry = linearizer.StateIndex().at(keys.back().GetLcmType());
+  return last_index_entry.offset + last_index_entry.tangent_dim;
+}
+
+}  // namespace internal
+
 template <typename ScalarType, typename NonlinearSolverType>
 void Optimizer<ScalarType, NonlinearSolverType>::ComputeAllCovariances(
     const Linearization<Scalar>& linearization,
@@ -171,18 +224,18 @@ template <typename ScalarType, typename NonlinearSolverType>
 void Optimizer<ScalarType, NonlinearSolverType>::ComputeCovariances(
     const Linearization<Scalar>& linearization, const std::vector<Key>& keys,
     std::unordered_map<Key, MatrixX<Scalar>>& covariances_by_key) {
-  size_t block_dim{};
-  const bool contiguous = linearizer_.CheckKeysAreContiguousAtStart(keys, &block_dim);
-  SYM_ASSERT(contiguous);
+  const bool same_order = internal::CheckKeyOrderMatchesLinearizerKeysStart(linearizer_, keys);
+  SYM_ASSERT(same_order);
+  const size_t block_dim = internal::ComputeBlockDimension(linearizer_, keys);
 
   // Copy into modifiable storage
   compute_covariances_storage_.H_damped = linearization.hessian_lower;
 
   internal::ComputeCovarianceBlockWithSchurComplement(compute_covariances_storage_.H_damped,
                                                       block_dim, epsilon_,
                                                       compute_covariances_storage_.covariance);
-  linearizer_.SplitCovariancesByKey(compute_covariances_storage_.covariance, keys,
-                                    covariances_by_key);
+  internal::SplitCovariancesByKey(linearizer_, compute_covariances_storage_.covariance, keys,
+                                  covariances_by_key);
 }
 
 template <typename ScalarType, typename NonlinearSolverType>