OptimizerLMBatch.hpp

Go to the documentation of this file.

/*
Copyright 2024 by Julian Nubert, Robotic Systems Lab, ETH Zurich.
All rights reserved.
This file is released under the "BSD-3-Clause License".
Please see the LICENSE file that has been included as part of this package.
 */
 
#ifndef OPTIMIZER_LM_BATCH_HPP
#define OPTIMIZER_LM_BATCH_HPP
 
// GTSAM
#include <gtsam/nonlinear/Marginals.h>
 
// Workspace
#include <holistic_fusion/core/optimizer/OptimizerLM.hpp>
 
namespace holistic_fusion {
 
class OptimizerLMBatch : public OptimizerLM {
 public:
  explicit OptimizerLMBatch(const std::shared_ptr<GraphConfig> graphConfigPtr) : OptimizerLM(graphConfigPtr) {
    // Initialize Slow Bundle Adjustement Smoother (if desired) -----------------------------------------------
    if (graphConfigPtr_->useAdditionalSlowBatchSmootherFlag_) {
      REGULAR_COUT << YELLOW_START << "HolisticFusion: OptimizerLMBatch" << GREEN_START
                   << " Initializing slow batch smoother that is optimized with LM." << COLOR_END << std::endl;
      // Set whether to use Cholesky factorization
      if (graphConfigPtr_->slowBatchSmootherUseCholeskyFactorizationFlag_) {
        lmParams_.linearSolverType = gtsam::NonlinearOptimizerParams::MULTIFRONTAL_CHOLESKY;
        REGULAR_COUT << "Using Multifrontal-Cholesky factorization for slow batch smoother (LM)." << std::endl;
      } else {
        lmParams_.linearSolverType = gtsam::NonlinearOptimizerParams::MULTIFRONTAL_QR;
        REGULAR_COUT << "Using Multifrontal-QR factorization for slow batch smoother (LM)." << std::endl;
      }
      // Print
      lmParams_.print("HolisticFusion: OptimizerLMBatch, LM Parameters:");
    }
  }
  ~OptimizerLMBatch() = default;
 
  // Empty update call --> do nothing
  bool update() override { return true; }
 
  // Add to graph without running optimization
  bool update(const gtsam::NonlinearFactorGraph& newGraphFactors, const gtsam::Values& newGraphValues,
              const std::map<gtsam::Key, double>& newGraphKeysTimeStampMap, const int depth = 0) override {
    // Add Bundle Adjustement Smoother factors to batch without running optimization
    containerBatchSmootherFactors_.add(newGraphFactors);
    containerBatchSmootherValues_.insert(newGraphValues);
    // Add to keyTimestampMap
    batchSmootherKeyTimestampMap_.insert(newGraphKeysTimeStampMap.begin(), newGraphKeysTimeStampMap.end());
    return true;
  }
 
  // Implementation
  bool updateExistingValues(const gtsam::Values& newGraphValues) override {
    // Update Values
    for (const auto& key : newGraphValues.keys()) {
      // Check whether key is in the batchSmootherValues
      if (containerBatchSmootherValues_.exists(key)) {
        // Remove old value
        containerBatchSmootherValues_.erase(key);
        // Insert new value
        containerBatchSmootherValues_.insert(key, newGraphValues.at(key));
        //        std::cout << "HolisticFusion: OptimizerLMBatch: updateExistingValues: Key " << gtsam::Symbol(key) << " updated." << std::endl;
      } else {
        throw std::logic_error("HolisticFusion: OptimizerLMBatch: updateExistingValues: Key " + gtsam::Symbol(key).string() +
                               " does not exist in batchSmootherValues.");
      }
    }
    return true;
  }
 
  // Optimize
  void optimize(int maxIterations) override {
    // Print
    std::cout << YELLOW_START << "HolisticFusion: OptimizerIsam2Batch" << GREEN_START << " Optimizing slow batch smoother." << COLOR_END
              << std::endl;
 
    // Set LM Parameters
    lmParams_.maxIterations = maxIterations;
 
    // Initialize Slow Bundle Adjustment Smoother
    batchSmootherPtr_ =
        std::make_shared<gtsam::LevenbergMarquardtOptimizer>(containerBatchSmootherFactors_, containerBatchSmootherValues_, lmParams_);
 
    // Do not reset containers --> Will build up graph again from scratch for next optimization
    // containerBatchSmootherFactors_.resize(0);
    // containerBatchSmootherValues_.clear();
    // Log State of graph in order to compute marginal covariance if desired
    graphLastOptimizedResult_ = containerBatchSmootherFactors_;
 
    // Optimize
    valuesLastOptimizedResult_ = batchSmootherPtr_->optimize();
    marginalsComputedForLastOptimizedResultFlag_ = false;
    optimizedAtLeastOnceFlag_ = true;
 
    // Set Sorted Keys by Time Flag
    dividedIntoSubGraphsFlag_ = false;
  }
 
  // Optimize bundle adjustement smoother (if desired)
  const gtsam::Values& getAllOptimizedStates() override {
    // Check
    if (!optimizedAtLeastOnceFlag_) {
      throw std::runtime_error("HolisticFusion: OptimizerLMBatch: getAllOptimizedStates: No optimization has been performed yet.");
    }
 
    // Return result
    return valuesLastOptimizedResult_;
  }
 
  // Get all keys of optimized states
  gtsam::KeyVector getAllOptimizedKeys() override { return valuesLastOptimizedResult_.keys(); }
 
  // Get nonlinear factor graph
  [[nodiscard]] const gtsam::NonlinearFactorGraph& getNonlinearFactorGraph() const override { return batchSmootherPtr_->graph(); }
 
  // Get keyTimestampMap
  const std::map<gtsam::Key, double>& getFullKeyTimestampMap() override { return batchSmootherKeyTimestampMap_; }
 
  // Calculate State / Covariance --------------------------------------------------------------
  // Pose3
  gtsam::Pose3 calculateEstimatedPose3(const gtsam::Key& key) override {
    if (!optimizedAtLeastOnceFlag_) {
      throw std::runtime_error("HolisticFusion: OptimizerLMBatch: calculateEstimatedPose: No optimization has been performed yet.");
    }
    return valuesLastOptimizedResult_.at<gtsam::Pose3>(key);
  }
  // Velocity3
  gtsam::Vector3 calculateEstimatedVelocity3(const gtsam::Key& key) override {
    if (!optimizedAtLeastOnceFlag_) {
      throw std::runtime_error("HolisticFusion: OptimizerLMBatch: calculateEstimatedVelocity: No optimization has been performed yet.");
    }
    return valuesLastOptimizedResult_.at<gtsam::Vector3>(key);
  }
  // Bias
  gtsam::imuBias::ConstantBias calculateEstimatedBias(const gtsam::Key& key) override {
    if (!optimizedAtLeastOnceFlag_) {
      throw std::runtime_error("HolisticFusion: OptimizerLMBatch: calculateEstimatedBias: No optimization has been performed yet.");
    }
    return valuesLastOptimizedResult_.at<gtsam::imuBias::ConstantBias>(key);
  }
  // Point3
  gtsam::Point3 calculateEstimatedPoint3(const gtsam::Key& key) override {
    if (!optimizedAtLeastOnceFlag_) {
      throw std::runtime_error("HolisticFusion: OptimizerLMBatch: calculateEstimatedDisplacement: No optimization has been performed yet.");
    }
    return valuesLastOptimizedResult_.at<gtsam::Point3>(key);
  }
  // Vector
  gtsam::Vector calculateEstimatedVector(const gtsam::Key& key) override {
    if (!optimizedAtLeastOnceFlag_) {
      throw std::runtime_error("HolisticFusion: OptimizerLMBatch: calculateStateAtKey: No optimization has been performed yet.");
    }
    return valuesLastOptimizedResult_.at<gtsam::Vector>(key);
  }
 
  static std::vector<double> linspace(double start, double end, int num) {
    std::vector<double> result;
    if (num <= 0) return result;  // Empty vector if num is 0 or negative
    if (num == 1) {               // Single point case
      result.push_back(start);
      return result;
    }
 
    double step = (end - start) / (num - 1);
    for (int i = 0; i < num; ++i) {
      result.push_back(start + i * step);
    }
 
    return result;
  }
 
  static int filterFactorsConnectedToKey(gtsam::NonlinearFactorGraph& graph, const gtsam::Key& keyToFilter) {
    gtsam::NonlinearFactorGraph filteredGraph;
    int numFilteredFactors = 0;
    for (const auto& factor : graph) {
      bool foundKeyInFactor = false;
      for (const gtsam::Key& key : factor->keys()) {
        if (key == keyToFilter) {
          foundKeyInFactor = true;
          break;
        }
      }
      if (!foundKeyInFactor) {
        filteredGraph.add(factor);
      } else {
        numFilteredFactors++;
      }
    }
    graph = filteredGraph;
    return numFilteredFactors;
  }
 
  static gtsam::Marginals computeMarginals(gtsam::NonlinearFactorGraph& graph, gtsam::Values& values) {
    //    // Create linearized factor graph
    //    gtsam::GaussianFactorGraph linearizedGraph = *graph.linearize(values);
 
    // Try to create marginals, filter out ill-conditioned states if necessary
    gtsam::Marginals marginals;
    try {
      marginals = gtsam::Marginals(graph, values);
    } catch (const gtsam::IndeterminateSystemException& e) {
      std::cout << "HolisticFusion: OptimizerLMBatch: computing marginals failed with error: " << e.what() << std::endl;
      // Filter out ill-conditioned states
      std::cout << "HolisticFusion: OptimizerLMBatch: filtering out ill-conditioned state: " << gtsam::Symbol(e.nearbyVariable()) << std::endl;
      // Values
      bool foundInValues = false;
      for (const gtsam::Key& key : values.keys()) {
        if (key == e.nearbyVariable()) {
          values.erase(key);
          foundInValues = true;
          break;
        }
      }
      // Factors
      int numFilteredFactors = filterFactorsConnectedToKey(graph, e.nearbyVariable());
 
      // Filtered out
      if (foundInValues) {
        std::cout << "HolisticFusion: OptimizerLMBatch: Filtered out state from values (" << gtsam::Symbol(e.nearbyVariable()) << ")."
                  << std::endl;
        std::cout << "HolisticFusion: OptimizerLMBatch: Filtered out " << numFilteredFactors << " factors. Retrying to compute marginals."
                  << std::endl;
        // Retry
        marginals = OptimizerLMBatch::computeMarginals(graph, values);
      } else {
        throw std::runtime_error("HolisticFusion: OptimizerLMBatch: computing marginals failed and could not filter out ill-conditioned state.");
      }
    } catch (const gtsam::ValuesKeyDoesNotExist& e) {
      std::cout << "HolisticFusion: OptimizerLMBatch: computing marginals failed with error: " << e.what() << std::endl;
 
      // Filtering out the factors that contain the not existing key
      int numFilteredFactors = filterFactorsConnectedToKey(graph, e.key());
      std::cout << "HolisticFusion: OptimizerLMBatch: Filtered out " << numFilteredFactors << " factors. Retrying to compute marginals."
                << std::endl;
      // Retry
      marginals = OptimizerLMBatch::computeMarginals(graph, values);
    }
 
    // Return
    return marginals;
  }
 
  void divideGraphIntoSubGraphs() {
    // Check
    if (!optimizedAtLeastOnceFlag_) {
      throw std::runtime_error("HolisticFusion: OptimizerLMBatch: computeAndStoreAllMarginals: No optimization has been performed yet.");
    }
 
    // All Keys
    gtsam::KeyVector keysSorted = valuesLastOptimizedResult_.keys();
 
    // Sort keys by time (given in keyTimestampMap)
    std::sort(keysSorted.begin(), keysSorted.end(), [&](const gtsam::Key& key1, const gtsam::Key& key2) {
      return batchSmootherKeyTimestampMap_.at(key1) < batchSmootherKeyTimestampMap_.at(key2);
    });
 
    // Get time span and number of keys
    double startTime = batchSmootherKeyTimestampMap_.at(keysSorted.front());
    double endTime = batchSmootherKeyTimestampMap_.at(keysSorted.back());
    int numKeys = keysSorted.size();
    double timeDifference = endTime - startTime;
 
    // Print
    std::cout << "HolisticFusion: OptimizerLMBatch: Graph consisting of " << numKeys << " keys is spanning " << timeDifference << " seconds."
              << std::endl;
 
    // Determine how many overlapping sub-graphs of size windowSizeSecond to create
    const double& windowSizeSeconds = graphConfigPtr_->windowSizeSecondsForMarginalsComputation_;  // Alias
    int numSubGraphs = std::ceil(timeDifference / windowSizeSeconds);
    std::cout << "HolisticFusion: OptimizerLMBatch: Dividing graph into " << numSubGraphs << " main windows." << std::endl;
    numSubGraphs += (numSubGraphs - 1);  // Overlapping windows
    std::cout << "HolisticFusion: OptimizerLMBatch: Dividing graph into " << numSubGraphs
              << " overlapping windows (to make sure there are no edges)." << std::endl;
 
    // Placeholders
    std::vector<double> subGraphCenterTimes;
    // Sub-graph start and end times
    std::vector<double> subGraphStartTimes;
    std::vector<double> subGraphEndTimes;
 
    // More than one sub-graph
    if (numSubGraphs > 1) {
      // Sub-graph Center Times: linearly spaced from startTime + windowSizeSeconds/2 to endTime - windowSizeSeconds/2
      subGraphCenterTimes = linspace(startTime + windowSizeSeconds / 2, endTime - windowSizeSeconds / 2, numSubGraphs);
      for (int i = 0; i < numSubGraphs; ++i) {
        subGraphStartTimes.push_back(subGraphCenterTimes[i] - windowSizeSeconds / 2);
        subGraphEndTimes.push_back(subGraphCenterTimes[i] + windowSizeSeconds / 2);
      }
      // One sub-graph
    } else {
      // Only one sub-graph
      subGraphCenterTimes.push_back((startTime + endTime) / 2);
      subGraphStartTimes.push_back(startTime);
      subGraphEndTimes.push_back(endTime);
    }
 
    // Print Relative Times
    for (int i = 0; i < numSubGraphs; ++i) {
      std::cout << "Subgraph " << i << " starts at relative " << subGraphStartTimes[i] - startTime << " seconds and ends at relative "
                << subGraphEndTimes[i] - startTime << " seconds." << std::endl;
      std::cout << "Center time: " << subGraphCenterTimes[i] - startTime << " seconds." << std::endl;
    }
 
    // Reset all containers
    subGraphs_.clear();
    keyToSubGraphIndexMap_.clear();
    marginalsForSubGraphs_.clear();
 
    // Create sub-graphs
    subGraphs_.resize(numSubGraphs);
    keyToSubGraphIndexMap_.clear();
    marginalsForSubGraphs_.resize(numSubGraphs);
 
    // Go through all factors and add them to the sub-graph if they are within the time window
    for (const auto& factorPtr : graphLastOptimizedResult_) {
      bool factorAtLeastInOneSubGraph = false;
      // Check whether the factor contains any key within the window of all sub-graphs
      for (int i = 0; i < numSubGraphs; ++i) {
        // Go through all keys of factor
        int numReferenceFrameKeysInFactor = 0;
        for (const auto& factorKey : factorPtr->keys()) {
          // Skip keys of reference frame alignment factors for determining timestamp (as they can span a lot of time) unless it is random
          // walk factor with two reference frame keys
          if (gtsam::Symbol(factorKey).chr() == 'r') {
            numReferenceFrameKeysInFactor++;
            if (numReferenceFrameKeysInFactor == 1) {
              continue;  // Skip timestamp of first reference frame key
            } else if (numReferenceFrameKeysInFactor > 1) {
              std::cout << "Found random walk factor with (at least) two reference frame keys (including " << gtsam::Symbol(factorKey)
                        << ")." << std::endl;  // Do not skip
            }
          }
 
          // Check whether key is within window
          if (batchSmootherKeyTimestampMap_[factorKey] >= subGraphStartTimes[i] &&
              batchSmootherKeyTimestampMap_[factorKey] <= subGraphEndTimes[i]) {
            subGraphs_[i].add(factorPtr);
            factorAtLeastInOneSubGraph = true;
            break;
          }
        }
      }
      // Has to be in at least one sub-graph
      if (!factorAtLeastInOneSubGraph) {
        throw std::runtime_error("HolisticFusion: OptimizerLMBatch: Factor does not contain any key within the time window of any sub-graph.");
      }
    }
 
    // Assign each key to the sub-graph with the closest center time to know in which sub-graph to compute marginals
    for (const auto& key : keysSorted) {
      double keyTime = batchSmootherKeyTimestampMap_[key];
      // Find the closest center time
      double minTimeDifference = std::numeric_limits<double>::max();
      int closestSubGraphIndex = -1;
      // Go through all sub-graphs
      for (int i = 0; i < numSubGraphs; ++i) {
        const double timeDifference = std::abs(subGraphCenterTimes[i] - keyTime);
        if (timeDifference < minTimeDifference) {
          minTimeDifference = timeDifference;
          closestSubGraphIndex = i;
        }
      }
      // Assign key to sub-graph
      keyToSubGraphIndexMap_[key] = closestSubGraphIndex;
    }
 
    // Optimize each sub-graph and compute marginals
    for (int i = 0; i < numSubGraphs; ++i) {
      // Print sub-graph information
      std::cout << "HolisticFusion: OptimizerLMBatch: Optimizing subgraph " << i << " with " << subGraphs_[i].size() << " factors and "
                << subGraphs_[i].keys().size() << " keys." << std::endl;
      // Optimize
      marginalsForSubGraphs_[i] = computeMarginals(subGraphs_[i], valuesLastOptimizedResult_);
      // Print
      std::cout << "HolisticFusion: OptimizerLMBatch: Subgraph " << i << " optimized and marginals computed." << std::endl;
    }
 
    // Set flag
    dividedIntoSubGraphsFlag_ = true;
 
    // Print
    std::cout << "HolisticFusion: OptimizerLMBatch: Divided graph into subgraphs." << std::endl;
  }
 
  // Marginal Covariance
  gtsam::Matrix calculateMarginalCovarianceMatrixAtKey(const gtsam::Key& valueKey) override {
    if (!optimizedAtLeastOnceFlag_) {
      throw std::runtime_error("HolisticFusion: OptimizerLMBatch: marginalCovariance: No optimization has been performed yet.");
    }
 
    // Super slow check
    if constexpr (false) {
      // Check whether key exists in optimized result
      if (!valuesLastOptimizedResult_.exists(valueKey)) {
        std::cout << "Key " << gtsam::Symbol(valueKey) << " does not exist in optimized result." << std::endl;
        throw std::runtime_error("HolisticFusion: OptimizerLMBatch: marginalCovariance: Key does not exist in optimized result.");
      }
      // Check whether key exists in nonlinar factor graph
      if (!graphLastOptimizedResult_.keys().exists(valueKey)) {
        std::cout << "Key " << gtsam::Symbol(valueKey) << " does not exist in nonlinear factor graph." << std::endl;
        throw std::runtime_error("HolisticFusion: OptimizerLMBatch: marginalCovariance: Key does not exist in nonlinear factor graph.");
      }
    }
 
    // Only use window around key for marginal covariance ------------------------------------------------
    if (graphConfigPtr_->useWindowForMarginalsComputationFlag_) {
      if (!dividedIntoSubGraphsFlag_) {
        divideGraphIntoSubGraphs();
      }
      // Get sub-graph index
      const int& subGraphsIndex_ = keyToSubGraphIndexMap_[valueKey];  // Alias
      if (subGraphsIndex_ != currentOptimizedSubgraphIndex_) {
        std::cout << GREEN_START << "Switching to sub-graph " << subGraphsIndex_ << " for marginal covariance computation." << COLOR_END
                  << std::endl;
        currentOptimizedSubgraphIndex_ = subGraphsIndex_;
      }
      // Compute marginal covariance and measure time
      gtsam::Matrix marginalCovariance = marginalsForSubGraphs_[subGraphsIndex_].marginalCovariance(valueKey);
      return marginalCovariance;
    }
    // Normal case: Compute marginals for whole graph
    else {
      // Have to compute all marginals (if not done already for this result)
      if (!marginalsComputedForLastOptimizedResultFlag_) {
        std::cout << "Getting new marginals for last optimized result." << std::endl;
        marginalsForLastOptimizedResult_ = gtsam::Marginals(graphLastOptimizedResult_, valuesLastOptimizedResult_);
        marginalsComputedForLastOptimizedResultFlag_ = true;
      }
      // Return
      return marginalsForLastOptimizedResult_.marginalCovariance(valueKey);
    }
  }
 
 private:
  // Optimizer itself
  std::shared_ptr<gtsam::LevenbergMarquardtOptimizer> batchSmootherPtr_;
  // Containers for bundle adjustment smoother
  gtsam::NonlinearFactorGraph containerBatchSmootherFactors_;
  gtsam::Values containerBatchSmootherValues_;
  // Result and keyTimestampMap
  bool optimizedAtLeastOnceFlag_ = false;
  gtsam::Values valuesLastOptimizedResult_;
  std::map<gtsam::Key, double> batchSmootherKeyTimestampMap_;
  // Sub-graphs for marginal covariance computation
  bool dividedIntoSubGraphsFlag_ = false;
  std::vector<gtsam::NonlinearFactorGraph> subGraphs_;
  std::unordered_map<gtsam::Key, int> keyToSubGraphIndexMap_;
  std::vector<gtsam::Marginals> marginalsForSubGraphs_;
  // Container for ingredients of last optimization, e.g. for marginal covariance
  gtsam::NonlinearFactorGraph graphLastOptimizedResult_;
  bool marginalsComputedForLastOptimizedResultFlag_ = false;
  gtsam::Marginals marginalsForLastOptimizedResult_;
 
  // Indicator of current sub-graph
  int currentOptimizedSubgraphIndex_ = -1;
};
 
}  // namespace holistic_fusion
 
#endif  // OPTIMIZER_LM_BATCH_HPP