rbrott · September 12, 2025 05:53
diff --git a/handmade_linearizer.cc b/handmade_linearizer.cc
 // optimizer.h

 /* ----------------------------------------------------------------------------
 * SymForce - Copyright 2022, Skydio, Inc.
 * This source code is under the Apache 2.0 license found in the LICENSE file.
 * ---------------------------------------------------------------------------- */

 #pragma once

 #include <sym/util/epsilon.h>

 #include "./factor.h"
 #include "./internal/linearizer_selector.h"
 #include "./levenberg_marquardt_solver.h"
 #include "./linearization.h"
 #include "./optimization_stats.h"

 namespace sym {

 /**
 * Class for optimizing a nonlinear least-squares problem specified as a list of Factors.  For
 * efficient use, create once and call Optimize() multiple times with different initial
 * guesses, as long as the factors remain constant and the structure of the Values is identical.
 *
 * Not thread safe! Create one per thread.
 *
 * Example usage:
 *
 *     // Create a Values
 *     sym::Key key0{'R', 0};
 *     sym::Key key1{'R', 1};
 *     sym::Valuesd values;
 *     values.Set(key0, sym::Rot3d::Identity());
 *     values.Set(key1, sym::Rot3d::Identity());
 *
 *     // Create some factors
 *     std::vector<sym::Factord> factors;
 *     factors.push_back(sym::Factord::Jacobian(
 *         [epsilon](const sym::Rot3d& rot, Eigen::Vector3d* const res,
 *                   Eigen::Matrix3d* const jac) {
 *           const sym::Rot3d prior = sym::Rot3d::Random();
 *           const Eigen::Matrix3d sqrt_info = Eigen::Vector3d::Ones().asDiagonal();
 *           sym::PriorFactorRot3<double>(rot, prior, sqrt_info, epsilon, res, jac);
 *         },
 *         {key0}));
 *     factors.push_back(sym::Factord::Jacobian(
 *         [epsilon](const sym::Rot3d& rot, Eigen::Vector3d* const res,
 *                   Eigen::Matrix3d* const jac) {
 *           const sym::Rot3d prior = sym::Rot3d::Random();
 *           const Eigen::Matrix3d sqrt_info = Eigen::Vector3d::Ones().asDiagonal();
 *           sym::PriorFactorRot3<double>(rot, prior, sqrt_info, epsilon, res, jac);
 *         },
 *         {key1}));
 *     factors.push_back(sym::Factord::Jacobian(
 *         [epsilon](const sym::Rot3d& a, const sym::Rot3d& b, Eigen::Vector3d* const res,
 *                   Eigen::Matrix<double, 3, 6>* const jac) {
 *           const Eigen::Matrix3d sqrt_info = Eigen::Vector3d::Ones().asDiagonal();
 *           const sym::Rot3d a_T_b = sym::Rot3d::Random();
 *           sym::BetweenFactorRot3<double>(a, b, a_T_b, sqrt_info, epsilon, res, jac);
 *         },
 *         {key0, key1}));
 *
 *     // Set up the params
 *     sym::optimizer_params_t params = DefaultLmParams();
 *     params.iterations = 50;
 *     params.early_exit_min_reduction = 0.0001;
 *
 *     // Optimize
 *     sym::Optimizer<double> optimizer(params, factors, epsilon);
 *     optimizer.Optimize(values);
 *
 * See symforce/test/symforce_optimizer_test.cc for more examples
 */
 template <typename ScalarType, typename _NonlinearSolverType = LevenbergMarquardtSolver<ScalarType>,
  typename _LinearizerType = internal::LinearizerSelector_t<typename _NonlinearSolverType::MatrixType>>
 class Optimizer {
 public:
  using Scalar = ScalarType;
  using NonlinearSolverType = _NonlinearSolverType;
  using FailureReason = typename NonlinearSolverType::FailureReason;
  using MatrixType = typename NonlinearSolverType::MatrixType;
  using Stats = OptimizationStats<MatrixType>;
  using LinearizerType = _LinearizerType;

  /**
   * Base constructor
   *
   * @param params: The params to use for the optimizer and nonlinear solver
   * @param factors: The set of factors to include
   * @param name: The name of this optimizer to be used for printing debug information
   * @param keys: The set of keys to optimize.  If empty, will use all optimized keys touched by the
   *    factors
   * @param epsilon: Epsilon for numerical stability
   */
  Optimizer(const optimizer_params_t& params, std::vector<Factor<Scalar>> factors,
            const std::string& name = "sym::Optimize", std::vector<Key> keys = {},
            const Scalar epsilon = sym::kDefaultEpsilon<Scalar>);

  /**
   * Constructor that copies in factors and keys, with arguments for the nonlinear solver
   *
   * See the base constructor for argument descriptions, additional arguments are forwarded to the
   * constructor for NonlinearSolverType
   */
  template <typename... NonlinearSolverArgs>
  Optimizer(const optimizer_params_t& params, std::vector<Factor<Scalar>> factors,
            const std::string& name, std::vector<Key> keys, Scalar epsilon,
            NonlinearSolverArgs&&... nonlinear_solver_args);

  // This cannot be moved or copied because the linearization keeps a pointer to the factors
  Optimizer(Optimizer&&) = delete;
  Optimizer& operator=(Optimizer&&) = delete;
  Optimizer(const Optimizer&) = delete;
  Optimizer& operator=(const Optimizer&) = delete;

  virtual ~Optimizer() = default;

  /**
   * Optimize the given values in-place
   *
   * @param num_iterations: If < 0 (the default), uses the number of iterations specified by the
   *    params at construction
   * @param populate_best_linearization: If true, the linearization at the best
   *    values will be filled out in the stats
   *
   * @returns The optimization stats
   */
  Stats Optimize(Values<Scalar>& values, int num_iterations = -1,
                 bool populate_best_linearization = false);

  /**
   * Optimize the given values in-place
   *
   * This overload takes the stats as an argument, and stores into there.  This allows users to
   * avoid reallocating memory for any of the entries in the stats, for use cases where that's
   * important.
   *
   * @param num_iterations: If < 0 (the default), uses the number of iterations specified by the
   *    params at construction
   * @param populate_best_linearization: If true, the linearization at the best values will be
   *    filled out in the stats
   * @param stats: An OptimizationStats to fill out with the result - if filling out dynamically
   *    allocated fields here, will not reallocate if memory is already allocated in the required
   *    shape (e.g. for repeated calls to Optimize())
   */
  virtual void Optimize(Values<Scalar>& values, int num_iterations,
                        bool populate_best_linearization, Stats& stats);

  /**
   * Optimize the given values in-place
   *
   * This overload takes the stats as an argument, and stores into there.  This allows users to
   * avoid reallocating memory for any of the entries in the stats, for use cases where that's
   * important.
   *
   * @param num_iterations: If < 0 (the default), uses the number of iterations specified by the
   *    params at construction
   * @param stats: An OptimizationStats to fill out with the result - if filling out dynamically
   *    allocated fields here, will not reallocate if memory is already allocated in the
   *    required shape (e.g. for repeated calls to Optimize())
   */
  void Optimize(Values<Scalar>& values, int num_iterations, Stats& stats);

  /**
   * Optimize the given values in-place
   *
   * This overload takes the stats as an argument, and stores into there.  This allows users to
   * avoid reallocating memory for any of the entries in the stats, for use cases where that's
   * important.
   *
   * @param stats: An OptimizationStats to fill out with the result - if filling out dynamically
   *    allocated fields here, will not reallocate if memory is already allocated in the
   *    required shape (e.g. for repeated calls to Optimize())
   */
  void Optimize(Values<Scalar>& values, Stats& stats);

  /**
   * Linearize the problem around the given values
   */
  Linearization<MatrixType> Linearize(const Values<Scalar>& values);

  /**
   * Get covariances for each optimized key at the given linearization
   *
   * Will reuse entries in covariances_by_key, allocating new entries so that the result contains
   * exactly the set of keys optimized by this Optimizer.  `covariances_by_key` must not contain any
   * keys that are not optimized by this Optimizer.
   *
   * May not be called before either Optimize() or Linearize() has been called.
   */
  void ComputeAllCovariances(const Linearization<MatrixType>& linearization,
                             std::unordered_map<Key, MatrixX<Scalar>>& covariances_by_key);

  /**
   * Get covariances for the given subset of keys at the given linearization
   *
   * This version is potentially much more efficient in both time and memory than computing the
   * covariances for all keys in the problem, by using the Schur complement trick.
   *
   * Currently requires that `keys` corresponds to a set of keys at the start of the list of keys
   * for the full problem, and in the same order.
   *
   * Given a matrix of structure:
   *
   *     A = ( B    E )
   *         ( E^T  C )
   *
   * where B corresponds to the variables in @param keys;
   *
   * If @param c_is_block_diagonal is true, the C matrix must be block diagonal, and a more
   * efficient solver for that structure is used.
   *
   * If @param c_is_block_diagonal is false, C can have any structure, but the computation is more
   * efficient if C is easy to factorize.
   *
   * Will reuse entries in `covariances_by_key`, allocating new entries so that the result contains
   * exactly the set of keys requested.  `covariances_by_key` must not contain any keys that are not
   * in `keys`.
   */
  Eigen::ComputationInfo ComputeCovariances(
      const Linearization<MatrixType>& linearization, const std::vector<Key>& keys,
      std::unordered_map<Key, MatrixX<Scalar>>& covariances_by_key,
      bool c_is_block_diagonal = true);

  /**
   * Get the full problem covariance at the given linearization
   *
   * Unlike ComputeCovariance and ComputeAllCovariances, this includes the off-diagonal blocks, i.e.
   * the cross-covariances between different keys.
   *
   * The ordering of entries here is the same as the ordering of the keys in the linearization,
   * which can be accessed via ``optimizer.Linearizer().StateIndex()``.
   *
   * May not be called before either Optimize() or Linearize() has been called.
   *
   * @param covariance A matrix that will be filled out with the full problem covariance.
   */
  void ComputeFullCovariance(const Linearization<MatrixType>& linearization,
                             MatrixX<Scalar>& covariance);

  /**
   * Get the optimized keys
   */
  const std::vector<Key>& Keys() const;

  /**
   * Get the factors
   */
  const std::vector<Factor<Scalar>>& Factors() const;

  /**
   * Get the Linearizer object
   */
  const LinearizerType& Linearizer() const;
  LinearizerType& Linearizer();

  /**
   * Get the NonlinearSolver object
   */
  const NonlinearSolverType& NonlinearSolver() const;
  NonlinearSolverType& NonlinearSolver();

  /**
   * Update the optimizer params
   */
  void UpdateParams(const optimizer_params_t& params);

  /**
   * Get the params used by the nonlinear solver
   */
  const optimizer_params_t& Params() const;

 protected:
  /**
   * Build the `linearize_func` functor for the underlying nonlinear solver
   */
  typename NonlinearSolverType::LinearizeFunc BuildLinearizeFunc(bool check_derivatives);

  bool IsInitialized() const;

  /**
   * Do initialization that depends on having a values
   */
  void Initialize(const Values<Scalar>& values);

  /**
   * If verbose_, log the exit status of the optimization
   */
  void MaybeLogStatus(const Stats& stats) const;

  const std::string& GetName();

  /// Store a copy of the nonlinear factors. The Linearization object in the state keeps a
  /// pointer to this memory.
  std::vector<Factor<Scalar>> factors_;

  /// The name of this optimizer to be used for printing debug information.
  std::string name_;

  /// Underlying nonlinear solver class.
  NonlinearSolverType nonlinear_solver_;

  Scalar epsilon_;
  bool debug_stats_;
  bool include_jacobians_;

  std::vector<Key> keys_;
  index_t index_;

  LinearizerType linearizer_;

  /*
   * Covariance matrix and damped Hessian, only used by ComputeCovariances() but cached here to save
   * reallocations. This may be the full problem covariance, or a subblock; it's always the full
   * problem Hessian
   */
  struct ComputeCovariancesStorage {
    sym::MatrixX<Scalar> covariance;
    MatrixType H_damped;
  };

  mutable ComputeCovariancesStorage compute_covariances_storage_;

  /// Functor for interfacing with the optimizer
  typename NonlinearSolverType::LinearizeFunc linearize_func_;

  bool verbose_;
 };

 // Shorthand instantiations
 using Optimizerd = Optimizer<double>;
 using Optimizerf = Optimizer<float>;

 /**
 * Simple wrapper to make it one function call.
 */
 template <typename Scalar, typename NonlinearSolverType = LevenbergMarquardtSolver<Scalar>>
 typename Optimizer<Scalar, NonlinearSolverType>::Stats Optimize(
    const optimizer_params_t& params, const std::vector<Factor<Scalar>>& factors,
    Values<Scalar>& values, const Scalar epsilon = kDefaultEpsilon<Scalar>) {
  Optimizer<Scalar, NonlinearSolverType> optimizer(params, factors, "sym::Optimize", {}, epsilon);
  return optimizer.Optimize(values);
 }

 /**
 * Sensible default parameters for Optimizer
 */
 optimizer_params_t DefaultOptimizerParams();

 }  // namespace sym

 #include "./optimizer.tcc"

 // Explicit instantiation declaration
 extern template class sym::Optimizer<double>;
 extern template class sym::Optimizer<float>;


 // optimizer.tcc
 /* ----------------------------------------------------------------------------
 * SymForce - Copyright 2022, Skydio, Inc.
 * This source code is under the Apache 2.0 license found in the LICENSE file.
 * ---------------------------------------------------------------------------- */

 #pragma once

 #include "./assert.h"
 #include "./internal/covariance_utils.h"
 #include "./internal/derivative_checker.h"
 #include "./internal/optimizer_utils.h"
 #include "./optimizer.h"

 namespace sym {

 // ----------------------------------------------------------------------------
 // Constructors
 // ----------------------------------------------------------------------------

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Optimizer(const optimizer_params_t& params,
                                                                      std::vector<Factor<Scalar>> factors,
                                                                      const std::string& name,
                                                                      std::vector<Key> keys, const Scalar epsilon)
    : factors_(std::move(factors)),
      name_(name),
      nonlinear_solver_(params, name, epsilon),
      epsilon_(epsilon),
      debug_stats_(params.debug_stats),
      include_jacobians_(params.include_jacobians),
      keys_(keys.empty() ? ComputeKeysToOptimize(factors_) : std::move(keys)),
      index_(),
      linearizer_(name_, factors_, keys_, params.include_jacobians, params.debug_checks),
      linearize_func_(BuildLinearizeFunc(params.check_derivatives)),
      verbose_(params.verbose) {
  SYM_ASSERT(factors_.size() > 0);
  SYM_ASSERT(keys_.size() > 0);

  SYM_ASSERT(!params.check_derivatives || params.include_jacobians);
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 template <typename... NonlinearSolverArgs>
 Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Optimizer(
    const optimizer_params_t& params, std::vector<Factor<Scalar>> factors, const std::string& name,
    std::vector<Key> keys, Scalar epsilon, NonlinearSolverArgs&&... nonlinear_solver_args)
    : factors_(std::move(factors)),
      name_(name),
      nonlinear_solver_(params, name, epsilon,
                        std::forward<NonlinearSolverArgs>(nonlinear_solver_args)...),
      epsilon_(epsilon),
      debug_stats_(params.debug_stats),
      include_jacobians_(params.include_jacobians),
      keys_(keys.empty() ? ComputeKeysToOptimize(factors_) : std::move(keys)),
      index_(),
      linearizer_(name_, factors_, keys_, params.include_jacobians),
      linearize_func_(BuildLinearizeFunc(params.check_derivatives)),
      verbose_(params.verbose) {
  SYM_ASSERT(factors_.size() > 0);
  SYM_ASSERT(keys_.size() > 0);

  SYM_ASSERT(!params.check_derivatives || params.include_jacobians);
 }

 // ----------------------------------------------------------------------------
 // Public methods
 // ----------------------------------------------------------------------------

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 typename sym::Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Stats
 Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Optimize(Values<Scalar>& values, int num_iterations,
                                                                     bool populate_best_linearization) {
  Stats stats{};
  Optimize(values, num_iterations, populate_best_linearization, stats);
  return stats;
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 void Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Optimize(Values<Scalar>& values,
                                                          int num_iterations,
                                                          bool populate_best_linearization,
                                                          Stats& stats) {
  // Create the index for the values
  Initialize(values);

  // Call the static helper function to run the optimization
  OptimizeImpl(values, nonlinear_solver_, linearize_func_, num_iterations,
               populate_best_linearization, name_, verbose_, stats);
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 void Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Optimize(Values<Scalar>& values,
                                                          int num_iterations, Stats& stats) {
  return Optimize(values, num_iterations, false, stats);
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 void Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Optimize(Values<Scalar>& values, Stats& stats) {
  return Optimize(values, -1, false, stats);
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 Linearization<typename NonlinearSolverType::MatrixType>
 Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Linearize(const Values<Scalar>& values) {
  Initialize(values);

  Linearization<MatrixType> linearization;
  linearize_func_(values, linearization);
  return linearization;
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 void Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::ComputeAllCovariances(
    const Linearization<MatrixType>& linearization,
    std::unordered_map<Key, MatrixX<Scalar>>& covariances_by_key) {
  SYM_ASSERT(IsInitialized());
  nonlinear_solver_.ComputeCovariance(linearization.hessian_lower,
                                      compute_covariances_storage_.covariance);
  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 LinearizerType>
 bool CheckKeyOrderMatchesLinearizerKeysStart(const LinearizerType& 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 LinearizerType>
 size_t ComputeBlockDimension(const LinearizerType& 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, typename LinearizerType>
 Eigen::ComputationInfo Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::ComputeCovariances(
    const Linearization<MatrixType>& linearization, const std::vector<Key>& keys,
    std::unordered_map<Key, MatrixX<Scalar>>& covariances_by_key, const bool c_is_block_diagonal) {
  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;

  const auto info = internal::ComputeCovarianceBlockWithSchurComplement(
      compute_covariances_storage_.H_damped, block_dim, epsilon_,
      compute_covariances_storage_.covariance, c_is_block_diagonal);

  if (info != Eigen::Success) {
    return info;
  }

  internal::SplitCovariancesByKey(linearizer_, compute_covariances_storage_.covariance, keys,
                                  covariances_by_key);

  return Eigen::Success;
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 void Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::ComputeFullCovariance(
    const Linearization<MatrixType>& linearization, MatrixX<Scalar>& covariance) {
  SYM_ASSERT(IsInitialized());
  nonlinear_solver_.ComputeCovariance(linearization.hessian_lower, covariance);
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 const std::vector<Key>& Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Keys() const {
  return keys_;
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 const std::vector<Factor<ScalarType>>& Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Factors() const {
  return factors_;
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 const LinearizerType&
 Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Linearizer() const {
  return linearizer_;
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 LinearizerType&
 Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Linearizer() {
  return linearizer_;
 }

 template <typename ScalarType, typename _NonlinearSolverType, typename _LinearizerType>
 _NonlinearSolverType& Optimizer<ScalarType, _NonlinearSolverType, _LinearizerType>::NonlinearSolver() {
  return nonlinear_solver_;
 };

 template <typename ScalarType, typename _NonlinearSolverType, typename _LinearizerType>
 const _NonlinearSolverType& Optimizer<ScalarType, _NonlinearSolverType, _LinearizerType>::NonlinearSolver() const {
  return nonlinear_solver_;
 };

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 void Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::UpdateParams(const optimizer_params_t& params) {
  nonlinear_solver_.UpdateParams(params);
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 const optimizer_params_t& Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Params() const {
  return nonlinear_solver_.Params();
 }

 // ----------------------------------------------------------------------------
 // Protected methods
 // ----------------------------------------------------------------------------

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 bool Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::IsInitialized() const {
  return index_.entries.size() != 0;
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 typename NonlinearSolverType::LinearizeFunc
 Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::BuildLinearizeFunc(const bool check_derivatives) {
  return [this, check_derivatives](const Values<Scalar>& values,
                                   Linearization<MatrixType>& linearization) {
    linearizer_.Relinearize(values, linearization);

    if (check_derivatives) {
      SYM_ASSERT(internal::CheckDerivatives(linearizer_, values, index_, linearization, epsilon_));
    }
  };
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 void Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::Initialize(const Values<Scalar>& values) {
  if (!IsInitialized()) {
    index_ = values.CreateIndex(keys_);
    nonlinear_solver_.SetIndex(index_);
  }
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 void Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::MaybeLogStatus(const Stats& stats) const {
  if (verbose_) {
    LogStatus<Stats, NonlinearSolverType>(name_, stats);
  }
 }

 template <typename ScalarType, typename NonlinearSolverType, typename LinearizerType>
 const std::string& Optimizer<ScalarType, NonlinearSolverType, LinearizerType>::GetName() {
  return name_;
 }

 extern template class Optimizer<double>;
 extern template class Optimizer<float>;

 }  // namespace sym


 // handmade_linearizer.h
 /* ----------------------------------------------------------------------------
 * SymForce - Copyright 2022, Skydio, Inc.
 * This source code is under the Apache 2.0 license found in the LICENSE file.
 * ---------------------------------------------------------------------------- */

 #pragma once

 #include <string>
 #include <unordered_map>
 #include <vector>

 #include "./factor.h"
 #include "./linearization.h"
 #include "./internal/linearized_dense_factor_pool.h"

 #include <handmade-symforce/src/types.h>
 #include <handmade-symforce/src/arena.h>
 #include <handmade-symforce/src/alloc.h>
 #include <handmade-symforce/src/linearizer.h>

 namespace sym {

 class HandmadeLinearizer {
 public:
  // Only supports doubles for now
  using Scalar = double;

  using LinearizedDenseFactor = typename Factor<Scalar>::LinearizedDenseFactor; // todo: what is this for?
  using LinearizationType = SparseLinearization<Scalar>;

  /**
   * Construct a Linearizer from factors and optional keys
   *
   * @param name: name of linearizer used for debug information
   * @param factors: Only stores a pointer, MUST be in scope for the lifetime of this object!
   * @param key_order: If provided, acts as an ordered set of keys that form the state vector to
   *     optimize. Can equal the set of all factor keys or a subset of all factor keys. If not
   *     provided, it is computed from all keys for all factors using a default ordering.
   * @param debug_checks: Whether to perform additional sanity checks for NaNs.  This uses
   *    additional compute but not additional memory except for logging.
   */
  HandmadeLinearizer(const std::string& name, const std::vector<Factor<Scalar>>& factors,
                  const std::vector<Key>& key_order = {},
                  bool include_jacobians = false, bool debug_checks = false);

  ~HandmadeLinearizer();

  /**
   * Returns whether Relinearize() has already been called once.
   *
   * Matters because many calculations need to be called only on the first linearization that are
   * then cached for subsequent use. Also, if Relinearize() has already been called, then the
   * matrices in the linearization are expected to already be allocated to the right size.
   */
  bool IsInitialized() const;

  /**
   * The keys to optimize, in the order they appear in the state vector (i.e., in rhs).
   */
  const std::vector<Key>& Keys() const;

  /**
   * A map from all optimized keys in the problem to an index entry for the corresponding optimized
   * values (where the offset is into the problem state vector, i.e., the rhs of a linearization).
   *
   * Only read if IsInitialized() returns true (i.e., after the first call to Relinearize).
   */
  const std::unordered_map<key_t, index_entry_t>& StateIndex() const;

  /**
   * Update linearization at a new evaluation point.
   * This is more efficient than reconstructing this object repeatedly. On the first call, it will
   * allocate memory and perform analysis needed for efficient repeated relinearization.
   *
   * On the first call to Relinearize, the matrices in linearization will be allocated and sized
   * correctly. On subsequent calls, the matrices of linearization are expected to already be
   * allocated and sized correctly.
   *
   * TODO(aaron): This should be const except that it can initialize the object
   */
  void Relinearize(const Values<Scalar>& values, LinearizationType& linearization);

 private:
  // The name of this linearizer to be used for printing debug information.
  std::string name_;
  const std::vector<Factor<Scalar>>& factors_;
  std::vector<Key> keys_;
  std::unordered_map<key_t, index_entry_t> state_index_;
  bool is_initialized_{false};
  bool debug_checks_;

  u8* buf_;
  sym_arena arena_;
  sym_allocator alloc_;

  i32 nblocks_;
  std::vector<bool> flipped_;
  i32* Hl_block_rows_;
  i32* Hl_block_cols_;
  i32* Hl_block_nz_indices_;

  sym_linearization lin_;
  sym_linearizer lzr_;

  // Interned keys
  // TODO: I'm not even sure we need this outside of init
  std::unordered_map<sym::Key, i32> key_map_;

  std::vector<i32> factor_opt_key_offsets_;

  // The index for each factor in the values.  Cached the first time we linearize, to avoid repeated
  // unordered_map lookups
  std::vector<std::vector<index_entry_t>> factor_indices_;

  // Linearized factors - stores individual factor residuals, jacobians, etc
  internal::LinearizedDenseFactorPool<Scalar> linearized_dense_factor_pool_;  // one per Jacobian shape
 };

 }  // namespace sym


 // handmade_linearizer.cc
 /* ----------------------------------------------------------------------------
 * SymForce - Copyright 2022, Skydio, Inc.
 * This source code is under the Apache 2.0 license found in the LICENSE file.
 * ---------------------------------------------------------------------------- */

 #include "./handmade_linearizer.h"

 namespace sym {

 HandmadeLinearizer::HandmadeLinearizer(const std::string& name,
                                       const std::vector<Factor<Scalar>>& factors,
                                       const std::vector<Key>& key_order,
                                       const bool include_jacobians, const bool debug_checks)
    : name_(name),
      factors_{factors},
      state_index_{}, // todo: what goes here exactly?
      debug_checks_{debug_checks},
      buf_((u8*) malloc(1L << 30)), // 1 GiB seems fine
      arena_{
          .beg = buf_,
          .end = buf_ + (1L << 30),
      },
      alloc_{
        .malloc = sym_arena_malloc,
        .free = sym_arena_free,
        .ctx = &arena_
      }
      {
  SYM_ASSERT(!include_jacobians, "Jacobians are not supported yet");

  if (key_order.empty()) {
    keys_ = ComputeKeysToOptimize(factors);
  } else {
    keys_ = key_order;
  }

  // TODO: Support sparse factors
  for (const auto& factor : factors_) {
    SYM_ASSERT(!factor.IsSparse(), "Sparse factors are not supported yet");
  }
 }

 bool HandmadeLinearizer::IsInitialized() const {
  return is_initialized_;
 }

 const std::vector<Key>& HandmadeLinearizer::Keys() const {
  return keys_;
 }

 const std::unordered_map<key_t, index_entry_t>& HandmadeLinearizer::StateIndex() const {
  SYM_ASSERT(IsInitialized());
  return state_index_;
 }

 void HandmadeLinearizer::Relinearize(const Values<Scalar>& values,
                                              LinearizationType& linearization) {
  // Note: None of initialization depends on the data but it requires a values object for key sizes.
  if (!is_initialized_) {
    // Compute state vector index (including key sizes) and intern keys
    i32 nkeys = (i32) keys_.size();
    i32* key_sizes = (i32*) alloc_.malloc(nkeys * sizeof(i32), alloc_.ctx);
    int32_t offset = 0;
    for (size_t i = 0; i < keys_.size(); ++i) {
      const auto& key = keys_.at(i);

      auto entry = values.IndexEntryAt(key);
      entry.offset = offset;
      state_index_[key.GetLcmType()] = entry;
      offset += entry.tangent_dim;

      key_sizes[i] = entry.tangent_dim;

      key_map_[key] = (i32)i;
    }

    // Count un-deduped blocks in the Hessian (really triplets)
    nblocks_ = 0;
    factor_indices_.reserve(factors_.size());
    // Track these to make sure that all combined keys are touched by at least one factor.
    std::unordered_set<Key> keys_touched_by_factors;
    for (const auto& factor : factors_) {
      i32 noptkeys = factor.OptimizedKeys().size();
      nblocks_ += (noptkeys + 1) * noptkeys / 2;

      factor_indices_.push_back(values.CreateIndex(factor.AllKeys()).entries);
    }

    // Materialize triplets
    Hl_block_rows_ = (i32*) alloc_.malloc(nblocks_ * sizeof(i32), alloc_.ctx);
    Hl_block_cols_ = (i32*) alloc_.malloc(nblocks_ * sizeof(i32), alloc_.ctx);
    flipped_.resize(nblocks_);
    i32 block_index = 0;
    for (const auto& factor : factors_) {
      // TODO: key_indices could be reserved with a max value from a previous pass and then
      // never allocated in this loop
      std::vector<int> key_indices;
      key_indices.reserve(factor.OptimizedKeys().size());
      i32 tangent_offset = 0;
      for (const auto& key : factor.OptimizedKeys()) {
        const auto key_index = key_map_.at(key);
        key_indices.push_back(key_index);

        factor_opt_key_offsets_.push_back(tangent_offset);
        tangent_offset += key_sizes[key_index];

        keys_touched_by_factors.insert(key);
      }

      for (size_t i = 0; i < key_indices.size(); ++i) {
        i32 row_index = key_indices.at(i);
        for (size_t j = 0; j < i; ++j) {
          // rectangular
          if (key_indices.at(j) < row_index) {
            Hl_block_rows_[block_index] = row_index;
            Hl_block_cols_[block_index] = key_indices.at(j);
            flipped_.at(block_index) = false;
          } else {
            Hl_block_rows_[block_index] = key_indices.at(j);
            Hl_block_cols_[block_index] = row_index;
            flipped_.at(block_index) = true;
          }
          ++block_index;

        }

        // triangular
        Hl_block_rows_[block_index] = row_index;
        Hl_block_cols_[block_index] = row_index;
        ++block_index;
      }
    }
    if (keys_.size() != keys_touched_by_factors.size()) {
      for (const auto& key : keys_) {
        if (keys_touched_by_factors.count(key) == 0) {
          throw std::runtime_error(
              fmt::format("Key {} is in the state vector but is not optimized by any factor.", key));
        }
      }
    }

    // Create Hl block matrix
    Hl_block_nz_indices_ = (i32*) alloc_.malloc(nblocks_ * sizeof(i32), alloc_.ctx);
    sym_csc_mat Hl_block = sym_csc_from_pairs(Hl_block_rows_, Hl_block_cols_, nblocks_, nkeys, nkeys, Hl_block_nz_indices_, &alloc_);

    // TODO: Not sure if we want METIS here.
    i32* key_perm = (i32*) alloc_.malloc(nkeys * sizeof(i32), alloc_.ctx);
    for (i32 i = 0; i < nkeys; ++i) {
      key_perm[i] = i;
    }

    // Create linearizer, linearization structure
    lzr_ = sym_linearizer_new(
      Hl_block, Hl_block_nz_indices_, nblocks_,
      key_sizes, nkeys,
      key_perm,
      &lin_,
      &alloc_
    );

    // TODO: Is residual really needed? Not sure how to populate that without actually linearizing the factors...

    sym_csc_mat_free(Hl_block, &alloc_);

    alloc_.free(key_perm, nkeys * sizeof(i32), alloc_.ctx);
    alloc_.free(key_sizes, nkeys * sizeof(i32), alloc_.ctx);
  }

  linearization.rhs.resize(lin_.rhs.n);

  // TODO: Is there an easier way to do this?
  linearization.hessian_lower.resize(lin_.Hl.nrows, lin_.Hl.ncols);
  linearization.hessian_lower.resizeNonZeros(lin_.Hl.nnz);
  std::copy(lin_.Hl.col_starts, lin_.Hl.col_starts + lin_.Hl.ncols + 1, linearization.hessian_lower.outerIndexPtr());
  std::copy(lin_.Hl.row_indices, lin_.Hl.row_indices + lin_.Hl.nnz, linearization.hessian_lower.innerIndexPtr());
  SYM_ASSERT(linearization.hessian_lower.isCompressed());

  // TODO: is zeroing necessary here?
  linearization.rhs.setZero();
  // Careful! setZero() on a sparse matrix is not what you want unfortunately
  Eigen::Map<VectorX<Scalar>>(linearization.hessian_lower.valuePtr(),
                              linearization.hessian_lower.nonZeros())
      .setZero();

  lin_.rhs.data = linearization.rhs.data();
  lin_.Hl.data = linearization.hessian_lower.valuePtr();

  // Linearize and accumulate the factors
  i32 block_index = 0;

  // only needed for init but not easy to express
  // I suppose one could make a lambda/function eventually
  typename internal::LinearizedDenseFactorPool<Scalar>::SizeTracker dense_factor_size_tracker;
  LinearizedDenseFactor linearized_dense_factor_init{};
  i32 factor_opt_key_index_ = 0;
  std::vector<double> residual;
  for (size_t i = 0; i < factors_.size(); ++i) {
    auto& factor = factors_.at(i);
    SYM_ASSERT(!factor.IsSparse());

    // Linearize the factor
    // TODO: Is there a way to skip the Jacobian here? I think I worked on that at some point and abandoned it
    const auto& linearized_dense_factor = [&]() -> const LinearizedDenseFactor& {
      const auto& factor_index = factor_indices_.at(i);
      if (is_initialized_) {
        auto& linearized_dense_factor = linearized_dense_factor_pool_.at(i);
        factor.Linearize(values, linearized_dense_factor, &factor_index);
        if (debug_checks_) {
          internal::CheckLinearizedFactor(name_, factor, values, linearized_dense_factor, factor_index);
        }

        return linearized_dense_factor;
      } else {
        factor.Linearize(values, linearized_dense_factor_init, &factor_index);
        if (debug_checks_) {
          internal::CheckLinearizedFactor(name_, factor, values, linearized_dense_factor_init, factor_index);
        }

        linearized_dense_factor_pool_.AppendFactorSize(linearized_dense_factor_init.residual.rows(),
                                                  linearized_dense_factor_init.rhs.rows(),
                                                  dense_factor_size_tracker);

        return linearized_dense_factor_init;
      }
    }();

    // Update residual
    for (int res_i = 0; res_i < linearized_dense_factor.residual.size(); ++res_i) {
      residual.push_back(linearized_dense_factor.residual(res_i));
    }

    // std::cout << "factor rhs:\n" << static_cast<Eigen::MatrixXd>(linearized_dense_factor.rhs) << "\n";
    // std::cout << "factor Hl:\n" << static_cast<Eigen::MatrixXd>(linearized_dense_factor.hessian) << "\n";

    // Accumulate the factor linearization into the right spots
    // TODO: Is it worth storing this separately or asking the factor about this?
    // I don't think optimized keys will otherwise be accessed during the linearization process
    // Oh and maybe this should go before the compute heavy stuff? Not sure how the memory accesses should be arranged
    const auto noptkeys = factor.OptimizedKeys().size();
    for (size_t j = 0; j < noptkeys; ++j) {
      for (size_t k = 0; k < j; ++k) {
        // rectangular case
        i32 row_index = Hl_block_rows_[block_index];
        i32 col_index = Hl_block_cols_[block_index];
        sym_linearizer_add_hessian_rect_block(
          lzr_, lin_,
          block_index,
          flipped_.at(block_index) ? col_index : row_index,
          flipped_.at(block_index) ? row_index : col_index,
          // the C API doesn't use const
          const_cast<f64*>(linearized_dense_factor.hessian.data()),
          linearized_dense_factor.hessian.rows(),
          factor_opt_key_offsets_.at(factor_opt_key_index_ + j),
          factor_opt_key_offsets_.at(factor_opt_key_index_ + k)
        );
        ++block_index;
      }

      i32 index = Hl_block_rows_[block_index];
      i32 factor_opt_key_offset = factor_opt_key_offsets_.at(factor_opt_key_index_ + j);
      sym_linearizer_add_hessian_tri_block(
        lzr_, lin_,
        block_index, index,
        // the C API doesn't use const
        const_cast<f64*>(linearized_dense_factor.hessian.data()),
        linearized_dense_factor.hessian.rows(),
        factor_opt_key_offset
      );
      ++block_index;

      sym_linearizer_add_rhs_block(
        lzr_, lin_,
        index,
        const_cast<f64*>(linearized_dense_factor.rhs.data()),
        factor_opt_key_offset
      );
    }

    factor_opt_key_index_ += noptkeys;
  }

  // Populate residual
  linearization.residual.resize(residual.size());
  for (size_t i = 0; i < residual.size(); i++) {
    linearization.residual(i) = residual.at(i);
  }

  is_initialized_ = true;
 }

 HandmadeLinearizer::~HandmadeLinearizer() {
  if (is_initialized_) {
    alloc_.free(Hl_block_rows_, nblocks_ * sizeof(i32), alloc_.ctx);
    alloc_.free(Hl_block_cols_, nblocks_ * sizeof(i32), alloc_.ctx);

    alloc_.free(Hl_block_nz_indices_, nblocks_ * sizeof(i32), alloc_.ctx);

    sym_linearizer_free(lzr_, &alloc_);

    lin_.rhs.data = NULL;
    lin_.Hl.data = NULL;
    sym_linearization_free(lin_, &alloc_);
  }

  printf("nalloc = %d\n", arena_.nalloc);
  printf("max_nalloc = %d\n", arena_.max_nalloc);
  SYM_ASSERT(arena_.nalloc == 0);

  free(buf_);
 }

 }  // namespace sym
No results found