WnoaInterpFactor.h

Go to the documentation of this file.

/* ----------------------------------------------------------------------------
 
 * GTSAM Copyright 2010, Georgia Tech Research Corporation,
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)
 
 * See LICENSE for the license information
 
 * -------------------------------------------------------------------------- */

 
#pragma once
 
#include <gtsam/base/Lie.h>
#include <gtsam/base/Testable.h>
#include <gtsam/base/VectorSpace.h>
#include <gtsam/base/timing.h>
#include <gtsam/geometry/Point1.h>
#include <gtsam/geometry/Point2.h>
#include <gtsam/geometry/Point3.h>
#include <gtsam/geometry/Pose2.h>
#include <gtsam/geometry/Pose3.h>
#include <gtsam/inference/Key.h>
#include <gtsam/nonlinear/NonlinearFactor.h>
#include <gtsam/nonlinear/Values.h>
#include <gtsam/nonlinear/WnoaInterpolator.h>
#include <gtsam/nonlinear/WnoaStateData.h>
 
#include <array>
#include <stdexcept>
#include <unordered_set>
#include <utility>
#include <vector>
 
namespace gtsam {

template <class PoseType>
class WnoaInterpFactor : public NoiseModelFactor {
 private:
  using Base = NoiseModelFactor;
  using This = WnoaInterpFactor<PoseType>;
  using VelocityType = typename gtsam::traits<PoseType>::TangentVector;
  static constexpr int dim = traits<PoseType>::dimension;
 
  // Convenient matrices
  using Matrix2N = Eigen::Matrix<double, 2 * dim, 2 * dim>;
  using MatrixN = Eigen::Matrix<double, dim, dim>;
  using LambdaPsiMats = typename Interpolator<PoseType>::LambdaPsiMats;
 
  // Tell the compiler to import the base class's version of error
  // Note: this is required because we don't define error( const HybridValues& )
  // in this class, and without this line, the base class's error(const
  // HybridValues&) is hidden by the error(const Values&, OptionalMatrixVecType)
  // defined in this class.
  using Base::error;
 
  // Inner factor that is called on interpolated values
  const NoiseModelFactor::shared_ptr inner_factor_;
  // Interpolator object for the given PoseType
  const Interpolator<PoseType> interpolator_;
  // disable noise model updates
  const bool fixed_noise_model_;
  // map keys to interpolated state
  std::unordered_map<Key, StateData> key_to_interp_;
  // map interpolated state to border states.
  std::unordered_map<StateData, std::pair<StateData, StateData>>
      interp_to_borders_;
  // map outer key to outer key index (for Jacobians)
  std::unordered_map<Key, int> outer_key_to_index_;
  // map of precomputed matrices for interpolation, keyed by StateData
  std::unordered_map<StateData, std::shared_ptr<LambdaPsiMats>>
      lambda_psi_pre_comp_;
 
  // Cache inner key -> index mapping to avoid rebuilding in noise model calc
  std::unordered_map<Key, int> inner_key_to_index_;
 
  // This struct contains the indices of the contributing outer keys as well as
  // flags for whether the inner key is interpolated or not. Keeps track of
  // which estimated states contribute to a given interpolated key, and the
  // corresponding Jacobian blocks for efficient linearization.
  struct InnerKeyMapping {
    bool isInterpolated = false;
    // For non-interpolated keys
    int directOuterIndex = -1;
    // For interpolated keys: cached outer indices and keys
    int indexPoseLeft = -1, indexVelLeft = -1, indexPoseRight = -1,
        indexVelRight = -1;
    Key keyPoseLeft = 0, keyVelLeft = 0, keyPoseRight = 0, keyVelRight = 0;
  };
  std::vector<InnerKeyMapping> inner_key_mappings_;
 
 public:
  struct PassedInterpData {
    Values values;

    std::unordered_map<Key, std::array<Matrix, 4>> jacobians;

    std::unordered_map<StateData, Matrix2N> condCovs;
  };

  WnoaInterpFactor(const NoiseModelFactor::shared_ptr inner_factor,
                   const std::set<StateData> estimated_states,
                   const std::set<StateData> interp_states,
                   const Eigen::Matrix<double, dim, 1> q_psd_diag,
                   const bool fixed_noise_model = false,
                   const bool precomp_interp_mats = true)
      : Base(inner_factor->noiseModel()),
        inner_factor_(inner_factor),
        interpolator_(q_psd_diag),
        fixed_noise_model_(fixed_noise_model) {
    // PROCESS INTERPOLATED STATES
    // loop through interpolated states
    for (const StateData& state : interp_states) {
      // search for estimated state that upper bound current interpolated state
      // Note: lower_bound finds the least upper bound using time-based
      // comparator of StateData. We can use it to find the right border state
      // for interpolation.
      auto iter_est_state = estimated_states.lower_bound(state);
      // Check if right border state is out of bounds (i.e., interp time is
      // outside the range of estimated times)
      if (iter_est_state == estimated_states.begin()) {
        throw std::runtime_error(
            "Interpolated state time is before all estimated state times");
      } else if (iter_est_state == estimated_states.end()) {
        throw std::runtime_error(
            "Interpolated state time is after all estimated state times");
      } else {
        // decrement iterator (points to left border state)
        iter_est_state--;
        // map interpolated state to borderstates
        interp_to_borders_[state] =
            std::pair(*iter_est_state, *std::next(iter_est_state));
        // Keep track of the inner keys corresponding to a given interpolated
        // state for easy lookup when building outer keys and mapping jacobians
        // later
        key_to_interp_[state.pose] = state;
        key_to_interp_[state.velocity] = state;
      }
      // Precompute Lambda and Psi WNOA interpolation matrices
      if (precomp_interp_mats) {
        lambda_psi_pre_comp_[state] =
            std::make_shared<LambdaPsiMats>(interpolator_.getLambdaPsi(
                interp_to_borders_[state].first.time,
                interp_to_borders_[state].second.time, state.time));
      } else {
        lambda_psi_pre_comp_[state] = nullptr;
      }
    }
    // DEFINE KEYS
    // Define set of "outer" keys that this wrapper factor is defined on.
    std::unordered_set<Key> outer_key_set;
    for (Key key : inner_factor->keys()) {
      if (key_to_interp_.find(key) == key_to_interp_.end()) {
        // inner key is not interpolated, add to outer keys
        outer_key_set.insert(key);
      } else {
        // inner key is interpolated, add associated border state keys to this
        // factor's keys
        StateData& interp = key_to_interp_.at(key);          // get state
        auto [left, right] = interp_to_borders_.at(interp);  // get borders
        outer_key_set.insert(left.pose);                     // add border keys
        outer_key_set.insert(left.velocity);
        outer_key_set.insert(right.pose);
        outer_key_set.insert(right.velocity);
      }
    }
    // Convert to key vector (from set)
    keys_ = KeyVector(outer_key_set.begin(), outer_key_set.end());
 
    // map outer keys to their associated index (used when mapping jacobians
    // later)
    for (size_t i = 0; i < this->keys_.size(); i++) {
      outer_key_to_index_[this->keys_[i]] = i;
    }
    // Build inner key mappings once and cache inner key indices
    // We will use this mapping to know how to map inner Jacobian blocks to
    // outer
 
    // number of inner keys (keys associated with just the inner factor)
    const KeyVector& inner_keys_init = inner_factor_->keys();
 
    // Vector of structs that stores the mapping information for each inner key
    // (whether it's interpolated, and the corresponding outer key indices)
    inner_key_mappings_.resize(inner_keys_init.size());
 
    // This map allows us to quickly find the index of an inner key in the inner
    // factor's key ordering This is important for correctly mapping Jacobian
    // blocks during linearization. We build this map once in the constructor to
    // avoid redundant work later (slight speed up)
    inner_key_to_index_.reserve(inner_keys_init.size());
 
    // Loop through all inner keys and save the relevant mappings to outer keys
    for (size_t i = 0; i < inner_keys_init.size(); ++i) {
      // inner key
      Key innerKey = inner_keys_init[i];
      // current index
      inner_key_to_index_[innerKey] = static_cast<int>(i);
      InnerKeyMapping mapping;
      auto itInterp = key_to_interp_.find(innerKey);
 
      // If this condition is met, then this inner key is not interpolated and
      // directly corresponds to an outer key We can map the Jacobian block for
      // this inner key directly to the corresponding outer key index
      if (itInterp == key_to_interp_.end()) {
        auto itOuter = outer_key_to_index_.find(innerKey);
        if (itOuter != outer_key_to_index_.end())
          mapping.directOuterIndex = itOuter->second;
      } else {
        // Otherwise, this inner key is interpolated and we need to map it to
        // its corresponding border states The Jacobian blocks that connect them
        // to the inner key
        mapping.isInterpolated = true;
 
        // get border states for this interpolated key using the
        // interp_to_borders_ map we built earlier
        const StateData& sd = itInterp->second;
        const auto& br = interp_to_borders_.at(sd);
        const StateData& left = br.first;
        const StateData& right = br.second;
 
        // Map the border states to their corresponding outer key indices
        // This uses the outer_key_to_index_ map we built earlier
        // We will need these indices to know where to map Jacobian blocks
        // during linearization
        mapping.indexPoseLeft = outer_key_to_index_.at(left.pose);
        mapping.indexVelLeft = outer_key_to_index_.at(left.velocity);
        mapping.indexPoseRight = outer_key_to_index_.at(right.pose);
        mapping.indexVelRight = outer_key_to_index_.at(right.velocity);
 
        // also save the actual keys for clarity (might not be accessed during
        // compute)
        mapping.keyPoseLeft = left.pose;
        mapping.keyVelLeft = left.velocity;
        mapping.keyPoseRight = right.pose;
        mapping.keyVelRight = right.velocity;
      }
      inner_key_mappings_[i] = mapping;
    }
  };

  ~WnoaInterpFactor() override {};

  void print(
      const std::string& s = "",
      const KeyFormatter& keyFormatter = DefaultKeyFormatter) const override {
    std::cout << s << "WnoaInterpFactor on ";
    for (const auto& k : this->keys()) {
      std::cout << keyFormatter(k) << " ";  // raw numeric key
    }
    std::cout << std::endl;
    this->inner_factor_->print("Inner Factor: ");
  }

  bool equals(const NonlinearFactor& expected,
              double tol = 1e-9) const override {
    const This* e = dynamic_cast<const This*>(&expected);
    return e != nullptr && Base::equals(*e, tol);
  }

  using Base::unwhitenedError;

  Vector unwhitenedError(const Values& values,
                         OptionalMatrixVecType H = nullptr) const override {
    return computeInterpolatedError(values, H);
  }

  std::shared_ptr<GaussianFactor> linearize(const Values& x) const override {
    // Only linearize if the factor is active
    if (!active(x)) return std::shared_ptr<JacobianFactor>();
 
    // Compute residual and effective noise model.
    std::vector<Matrix> A(size());
    Vector b;
    auto noise_model = eval(x, nullptr, b, &A);
    return makeJacobianFactor(A, b, noise_model);
  }

  std::shared_ptr<GaussianFactor> linearize(
      const Values& x, PassedInterpData* passedInterpData) const {
    // Only linearize if the factor is active
    if (!active(x)) return std::shared_ptr<JacobianFactor>();
 
    // Compute residual and effective noise model.
    std::vector<Matrix> A(size());
    Vector b;
    auto noise_model = eval(x, passedInterpData, b, &A);
    return makeJacobianFactor(A, b, noise_model);
  }

  double error(const Values& c) const override {
    if (!active(c)) return 0.0;
 
    Vector b;
    auto noise_model = eval(c, nullptr, b, nullptr);
    return loss(b, noise_model);
  }

  double error(const Values& c, PassedInterpData* passedInterpData) const {
    if (!active(c)) return 0.0;
 
    Vector b;
    auto noise_model = eval(c, passedInterpData, b, nullptr);
    return loss(b, noise_model);
  }

  SharedGaussian noiseModel(Values& x) const {
    // if fixed noise then just return the standard gaussian model
    if (fixed_noise_model_) {
      return std::dynamic_pointer_cast<noiseModel::Gaussian>(
          Base::noiseModel());
    }
    // Call evaluate error to get inner Jacobians and convariances
    std::vector<Matrix> JacInner(inner_factor_->size());
    std::unordered_map<StateData, Matrix2N> InterpCondCovs;
    Vector b =
        -computeInterpolatedError(x, nullptr, &JacInner, &InterpCondCovs);
    // get interpolated noise model
    return getInterpolatedNoiseModel(JacInner, InterpCondCovs);
  }

  std::unordered_map<Key, StateData> getInterpolatedKeys() const {
    return key_to_interp_;
  }

  std::unordered_map<StateData, std::pair<StateData, StateData>>
  getInterpToBorders() const {
    return interp_to_borders_;
  }
 
 private:
  noiseModel::Gaussian::shared_ptr eval(const Values& values,
                                        PassedInterpData* passedInterpData,
                                        Vector& b,
                                        std::vector<Matrix>* A) const {
    if (A && (A->size() != size())) A->resize(size());
 
    return fixed_noise_model_ ? evalFixed(values, A, passedInterpData, b)
                              : evalInterp(values, A, passedInterpData, b);
  }

  noiseModel::Gaussian::shared_ptr evalFixed(const Values& values,
                                             std::vector<Matrix>* A,
                                             PassedInterpData* passedInterpData,
                                             Vector& b) const {
    b = -computeInterpolatedError(values, A, nullptr, nullptr,
                                  passedInterpData);
    return std::dynamic_pointer_cast<noiseModel::Gaussian>(Base::noiseModel());
  }

  noiseModel::Gaussian::shared_ptr evalInterp(
      const Values& values, std::vector<Matrix>* A,
      PassedInterpData* passedInterpData, Vector& b) const {
    std::vector<Matrix> jacInner(inner_factor_->size());
 
    if (passedInterpData) {
      b = -computeInterpolatedError(values, A, &jacInner, nullptr,
                                    passedInterpData);
      return getInterpolatedNoiseModel(jacInner, passedInterpData->condCovs);
    }
 
    std::unordered_map<StateData, Matrix2N> localInterpCondCovs;
    b = -computeInterpolatedError(values, A, &jacInner, &localInterpCondCovs);
    return getInterpolatedNoiseModel(jacInner, localInterpCondCovs);
  }

  std::shared_ptr<GaussianFactor> makeJacobianFactor(
      std::vector<Matrix>& A, Vector& b,
      const noiseModel::Gaussian::shared_ptr& noise_model) const {
    noise_model->WhitenSystem(A, b);
 
    std::vector<std::pair<Key, Matrix>> terms(size());
    for (size_t j = 0; j < size(); ++j) {
      terms[j].first = keys()[j];
      terms[j].second.swap(A[j]);
    }
 
    using noiseModel::Constrained;
    if (noiseModel_ && noiseModel_->isConstrained()) {
      return std::make_shared<JacobianFactor>(
          terms, b, std::static_pointer_cast<Constrained>(noiseModel_)->unit());
    }
    return std::make_shared<JacobianFactor>(terms, b);
  }

  double loss(const Vector& b,
              const noiseModel::Gaussian::shared_ptr& noise_model) const {
    if (noise_model)
      return noise_model->loss(noise_model->squaredMahalanobisDistance(b));
    return 0.5 * b.squaredNorm();
  }

  Vector computeInterpolatedError(
      const Values& values, OptionalMatrixVecType H = nullptr,
      OptionalMatrixVecType H_inner = nullptr,
      std::unordered_map<StateData, Matrix2N>* InterpCondCovs = nullptr,
      PassedInterpData* passedInterpData = nullptr) const {
    // Interpolation Jacobians stored as flattened map: per interpolated key ->
    // 4 blocks
    std::unordered_map<Key, std::array<Matrix, 4>> interpJacobiansLocal;
    Values valuesInterpLocal;
 
    std::unordered_map<Key, std::array<Matrix, 4>>* InterpJacobians = nullptr;
    Values* values_interp = nullptr;
 
    if (passedInterpData) {
      values_interp = &passedInterpData->values;
      if (H) InterpJacobians = &passedInterpData->jacobians;
      if (InterpCondCovs) {
        InterpCondCovs = &passedInterpData->condCovs;
      }
    } else {
      if (H) {
        InterpJacobians = &interpJacobiansLocal;
        valuesInterpLocal =
            getInterpolatedValues(values, InterpJacobians, InterpCondCovs);
        values_interp = &valuesInterpLocal;
      } else {
        valuesInterpLocal =
            getInterpolatedValues(values, nullptr, InterpCondCovs);
        values_interp = &valuesInterpLocal;
      }
    }
 
    // cache inner keys once
    const KeyVector& inner_keys = inner_factor_->keys();
 
    // construct values for inner factor using mappings
    Values values_inner;
    for (size_t i = 0; i < inner_keys.size(); ++i) {
      Key key = inner_keys[i];
      if (inner_key_mappings_[i].isInterpolated) {
        auto it_interp = values_interp->find(key);
        if (it_interp == values_interp->end())
          throw std::runtime_error("Interpolated key missing in values_interp");
        values_inner.insert(key, it_interp->value);
      } else {
        auto it_outer = values.find(key);
        if (it_outer == values.end())
          throw std::runtime_error("Key " + DefaultKeyFormatter(key) +
                                   " not found in outer values");
        values_inner.insert(key, it_outer->value);
      }
    }
 
    // Call inner factor error function with interpolated values.
    std::vector<Matrix> H_inner_local;
    Vector error;
    if (!H_inner) {
      // if H_inner not passed in, use local variable.
      H_inner_local.resize(inner_keys.size());
      H_inner = &H_inner_local;
    }
    if (H || !fixed_noise_model_) {
      error = inner_factor_->unwhitenedError(values_inner, H_inner);
    } else {
      error = inner_factor_->unwhitenedError(values_inner);
    }
 
    // compute Jacobians for outer keys
    if (H) {
      // loop through inner keys and update outer keys via backpropagation
      // NOTE: it is possible for two inner keys to affect the same outer key
      for (size_t i = 0; i < inner_keys.size(); i++) {
        const Key inner_key = inner_keys[i];
        const Matrix& Jinner = (*H_inner)[i];
        const InnerKeyMapping& mapping = inner_key_mappings_[i];
        if (mapping.isInterpolated) {
          const std::array<Matrix, 4>& J4 = InterpJacobians->at(inner_key);
          // Order: 0:LPose, 1:LVel, 2:RPose, 3:RVel
          if (mapping.indexPoseLeft >= 0) {
            const Matrix& Jblock = J4[0];
            if ((*H)[mapping.indexPoseLeft].size() == 0)
              (*H)[mapping.indexPoseLeft].setZero(Jinner.rows(), Jblock.cols());
            (*H)[mapping.indexPoseLeft].noalias() += Jinner * Jblock;
          }
          if (mapping.indexVelLeft >= 0) {
            const Matrix& Jblock = J4[1];
            if ((*H)[mapping.indexVelLeft].size() == 0)
              (*H)[mapping.indexVelLeft].setZero(Jinner.rows(), Jblock.cols());
            (*H)[mapping.indexVelLeft].noalias() += Jinner * Jblock;
          }
          if (mapping.indexPoseRight >= 0) {
            const Matrix& Jblock = J4[2];
            if ((*H)[mapping.indexPoseRight].size() == 0)
              (*H)[mapping.indexPoseRight].setZero(Jinner.rows(),
                                                   Jblock.cols());
            (*H)[mapping.indexPoseRight].noalias() += Jinner * Jblock;
          }
          if (mapping.indexVelRight >= 0) {
            const Matrix& Jblock = J4[3];
            if ((*H)[mapping.indexVelRight].size() == 0)
              (*H)[mapping.indexVelRight].setZero(Jinner.rows(), Jblock.cols());
            (*H)[mapping.indexVelRight].noalias() += Jinner * Jblock;
          }
        } else {
          const int k = mapping.directOuterIndex;
          if ((*H)[k].size() == 0)
            (*H)[k].setZero(Jinner.rows(), Jinner.cols());
          (*H)[k].noalias() += Jinner;
        }
      }
    }
 
    return error;
  }

  Values getInterpolatedValues(
      const Values& values,
      std::unordered_map<Key, std::array<Matrix, 4>>* InterpJacobians = nullptr,
      std::unordered_map<StateData, Matrix2N>* InterpCondCovs = nullptr) const {
    Values values_interp;  // interpolated values
 
    // loop through interpolated state map and compute values
    for (const auto& [interp_state, border_states] : interp_to_borders_) {
      // unpack border states
      auto& [left, right] = border_states;
      // retrieve estimated state values
      const auto state_left = TimestampedPoseVelocity<PoseType>(
          values.at<PoseType>(left.pose),
          values.at<VelocityType>(left.velocity), left.time);
 
      const auto state_right = TimestampedPoseVelocity<PoseType>(
          values.at<PoseType>(right.pose),
          values.at<VelocityType>(right.velocity), right.time);
 
      // Get interpolated state velocity pair
      PoseVelocity<PoseType> result;
 
      std::vector<Matrix> H(8);
      if (InterpJacobians) {
        result = interpolator_.interpolatePoseAndVelocity(
            state_left, state_right, interp_state.time, &H, nullptr, nullptr,
            lambda_psi_pre_comp_.at(interp_state));
      } else {
        result = interpolator_.interpolatePoseAndVelocity(
            state_left, state_right, interp_state.time, nullptr, nullptr,
            nullptr, lambda_psi_pre_comp_.at(interp_state));
      }
 
      // insert into values structure
      values_interp.insert(interp_state.pose, result.pose);
      values_interp.insert(interp_state.velocity, result.vel);
 
      // arrange jacobians in flattened map (fixed order blocks)
      if (InterpJacobians) {
        (*InterpJacobians)[interp_state.pose] =
            std::array<Matrix, 4>{H[0], H[1], H[2], H[3]};
        (*InterpJacobians)[interp_state.velocity] =
            std::array<Matrix, 4>{H[4], H[5], H[6], H[7]};
      }
 
      // Conditional covariance of interpolated states for noise model update
      if (InterpCondCovs) {
        auto state_tau =
            TimestampedPoseVelocity<PoseType>(result, interp_state.time);
        Matrix2N Sigma_tau = interpolator_.computeConditionalCov(
            state_left, state_right, state_tau);
        (*InterpCondCovs)[interp_state] =
            Sigma_tau;  // assumed preallocated vector
      }
    }
 
    return values_interp;
  }

  SharedGaussian getInterpolatedNoiseModel(
      const std::vector<Matrix>& Jacobians,
      const std::unordered_map<StateData, Matrix2N>& InterpCondCovs) const {
    // Get noise model of inner factor
    noiseModel::Gaussian::shared_ptr noise_model_ptr =
        std::dynamic_pointer_cast<noiseModel::Gaussian>(
            inner_factor_->noiseModel());
    // Check that the measurement noise is set up as a gaussian
    assert(noise_model_ptr &&
           "Noise model of inner factor must be noiseModel::Gaussian or "
           "derivative");
 
    // Initialize new covariance with the existing measurement covariance
    int err_dim = noise_model_ptr->dim();
    Matrix noise_cov = noise_model_ptr->covariance();
 
    // Use cached mapping from inner keys to indices for Jacobian lookup
    // Compute the covariance update based on interpolated states
    // Note: Here, we leverage the block-diagonal approximation of the
    // interpolated covariances (i.e., independence approximation)
    for (auto& [state, borders] : interp_to_borders_) {
      // Retrieve Jacobians from inner factor
      Matrix G_pose(err_dim, dim);
      Matrix G_vel(err_dim, dim);
      auto itPose = inner_key_to_index_.find(state.pose);
      if (itPose != inner_key_to_index_.end()) {
        G_pose = Jacobians[itPose->second];
      } else {
        G_pose.setZero();
      }
      auto itVel = inner_key_to_index_.find(state.velocity);
      if (itVel != inner_key_to_index_.end()) {
        G_vel = Jacobians[itVel->second];
      } else {
        G_vel.setZero();
      }
      Matrix G_tau(err_dim, 2 * dim);
      G_tau << G_pose, G_vel;
      // add covariance
 
      const Matrix2N& Sigma_tau = InterpCondCovs.at(state);
      noise_cov += G_tau * Sigma_tau * G_tau.transpose();
    }
 
    // Return interpolated noise model
    return noiseModel::Gaussian::Covariance(noise_cov);
  }
};

template <class POSE>
struct traits<WnoaInterpFactor<POSE>>
    : public Testable<WnoaInterpFactor<POSE>> {};
 
}  // namespace gtsam