BatchFactor.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/Testable.h>
#include <gtsam/linear/BatchJacobianFactor.h>
#include <gtsam/linear/HessianFactor.h>
#include <gtsam/linear/JacobianFactor.h>
#include <gtsam/linear/NoiseModel.h>
#include <gtsam/nonlinear/NonlinearFactor.h>
 
#include <Eigen/StdVector>
#include <algorithm>
#include <cassert>
#include <map>
#include <type_traits>
#include <vector>
 
namespace gtsam {
 
namespace detail {

template <typename FactorType, typename K1, typename K2, typename Meas,
          typename Model, typename... Args>
static FactorType createFactor(K1 k1, K2 k2, const Meas& z, const Model& model,
                               Args&&... args) {
  if constexpr (std::is_constructible_v<FactorType, K1, K2, Meas, Model,
                                        Args...>) {
    return FactorType(k1, k2, z, model, std::forward<Args>(args)...);
  } else if constexpr (std::is_constructible_v<FactorType, Meas, Model, K1, K2,
                                               Args...>) {
    return FactorType(z, model, k1, k2, std::forward<Args>(args)...);
  } else if constexpr (std::is_constructible_v<FactorType, K1, K2, Meas,
                                               Args..., Model>) {
    return FactorType(k1, k2, z, std::forward<Args>(args)..., model);
  } else {
    // This static_assert will trigger if none of the above match.
    // We repeat the check to produce a readable error message.
    static_assert(
        std::is_constructible_v<FactorType, K1, K2, Meas, Model, Args...>,
        "BatchFactor: Could not find a matching constructor for FactorType. "
        "Tried: (K1, K2, Z, Model, Args...), (Z, Model, K1, K2, Args...), (K1, "
        "K2, Z, Args..., Model)");
    return FactorType(k1, k2, z, model, std::forward<Args>(args)...);
  }
}
 
}  // namespace detail

template <typename FactorType, int ErrorDim>
class BatchFactor : public NonlinearFactor {
 public:
  // Static assertion to ensure FactorType derives from NoiseModelFactor
  static_assert(std::is_base_of<NoiseModelFactor, FactorType>::value,
                "FactorType must derive from NoiseModelFactor");
 
  using Base = NonlinearFactor;
  using This = BatchFactor<FactorType, ErrorDim>;
  using shared_ptr = std::shared_ptr<This>;
  using FactorVector =
      std::vector<FactorType, Eigen::aligned_allocator<FactorType>>;
 
 private:
  using Allocator = typename FactorVector::allocator_type;
  FactorVector factors_;  
  struct KeyInfo {
    Key key;
    int dim;
    size_t slot;
    DenseIndex offset;
  };
  std::vector<KeyInfo> keyInfo_;
  std::vector<std::array<DenseIndex, FactorType::N>> indices_;
  bool useHessianFactor_{false};
  bool allDimensionsFixed_{false};
  bool hasConstrainedNoiseModel_{false};
  bool allNoiseModelsAreUnit_{true};
  std::vector<size_t> keyDimensions_;
 
  template <size_t... Is>
  static constexpr bool hasFixedDimensions(std::index_sequence<Is...>) {
    return (
        (traits<typename FactorType::template ValueType<Is + 1>>::dimension !=
         Eigen::Dynamic) &&
        ...);
  }
 
  template <size_t... Is>
  size_t keyDimensionFromSlot(const Key key, size_t slot, const Values& values,
                             std::index_sequence<Is...>) const {
    const std::array<size_t, FactorType::N> dims = {
        traits<typename FactorType::template ValueType<Is + 1>>::GetDimension(
            values.at<typename FactorType::template ValueType<Is + 1>>(key))...};
 
    if (slot >= dims.size()) {
      throw std::runtime_error(
          "BatchFactor::keyDimensionFromSlot: slot out of range.");
    }
    const size_t dimension = dims[slot];
    return dimension;
  }
 
  const std::vector<size_t>& keyDimensions() const {
    if (!allDimensionsFixed_) {
      throw std::runtime_error(
          "BatchFactor::keyDimensions: cached dimensions unavailable for dynamic "
          "key types.");
    }
    return keyDimensions_;
  }
 
  std::vector<size_t> keyDimensions(const Values& values) const {
    std::vector<size_t> dims;
    dims.reserve(keyInfo_.size());
    for (const auto& info : keyInfo_) {
      if (info.dim >= 0) {
        dims.push_back(static_cast<size_t>(info.dim));
      } else {
        constexpr auto slots = std::make_index_sequence<FactorType::N>{};
        const size_t dim = keyDimensionFromSlot(info.key, info.slot, values, slots);
        if (dim == 0) {
          throw std::runtime_error(
              "BatchFactor::keyDimensions: cannot determine dynamic key "
              "dimension.");
        }
        dims.push_back(dim);
      }
    }
    return dims;
  }
 
  std::shared_ptr<JacobianFactor> linearizeToJacobianFactor(
      const Values& values) const {
    const size_t total_rows = factors_.size() * ErrorDim;
    const std::vector<size_t>& keyDims =
        allDimensionsFixed_ ? keyDimensions() : keyDimensions(values);
    VerticalBlockMatrix Ab(keyDims, total_rows, true);
    Ab.matrix().setZero();
 
    std::vector<Matrix> H(FactorType::N);
    Vector rowSigmas;
    Vector rowMus;
    if (hasConstrainedNoiseModel_) {
      rowSigmas = Vector::Ones(total_rows);
      rowMus = Vector::Constant(total_rows, 1000.0);
    }
    for (size_t i = 0; i < factors_.size(); ++i) {
      const auto& factor = factors_[i];
      const size_t row_start = i * ErrorDim;
      Vector rhs = -factor.unwhitenedError(values, H);
 
      const auto noise = factor.noiseModel();
      if (!allNoiseModelsAreUnit_ && noise && !noise->isUnit()) {
        factor.noiseModel()->WhitenSystem(H, rhs);
      }
 
      if (hasConstrainedNoiseModel_ && noise && noise->isConstrained()) {
        auto constrainedModel =
            std::dynamic_pointer_cast<noiseModel::Constrained>(noise);
        if (!constrainedModel) {
          throw std::runtime_error(
              "BatchFactor::linearizeToJacobianFactor: unsupported constrained "
              "noise model type.");
        }
        for (size_t r = 0; r < ErrorDim; ++r) {
          const size_t row = row_start + r;
          if (constrainedModel->constrained(r)) {
            rowMus[row] = constrainedModel->mu()[r];
            rowSigmas[row] = 0.0;
          } else {
            rowSigmas[row] = 1.0;
          }
        }
      }
 
      const auto& indices_i = indices_[i];
      for (size_t k = 0; k < FactorType::N; ++k) {
        const DenseIndex index = indices_i[k];
        Ab(index).block(row_start, 0, ErrorDim, H[k].cols()) = H[k];
      }
      Ab(keys().size()).block(row_start, 0, ErrorDim, 1) = rhs;
    }
 
    SharedDiagonal jacobianModel = noiseModel::Unit::Create(total_rows);
    if (hasConstrainedNoiseModel_) {
      jacobianModel = std::static_pointer_cast<noiseModel::Diagonal>(
          noiseModel::Constrained::MixedSigmas(rowMus, rowSigmas));
    }
    return std::make_shared<JacobianFactor>(
        keys(), std::move(Ab), jacobianModel);
  }
 
  template <size_t... Is>
  std::shared_ptr<GaussianFactor> linearizeToBatchJacobian(
      const Values& values, std::index_sequence<Is...>) const {
    (void)values;
    using CompactFactor = BatchJacobianFactor<
        ErrorDim,
        traits<typename FactorType::template ValueType<Is + 1>>::dimension...>;
    auto batch = std::make_shared<CompactFactor>(keys(), keyDimensions());
    batch->reserve(factors_.size());
 
    std::vector<Matrix> H(FactorType::N);
    for (size_t i = 0; i < factors_.size(); ++i) {
      const auto& factor = factors_[i];
      Vector rhs = -factor.unwhitenedError(values, H);
 
      if (factor.noiseModel() && !factor.noiseModel()->isUnit()) {
        factor.noiseModel()->WhitenSystem(H, rhs);
      }
 
      batch->addRow(indices_[i], H, rhs);
    }
    return batch;
  }
 
 public:

  BatchFactor() = default;

  const FactorVector& factors() const { return factors_; }

  size_t numFactors() const { return factors_.size(); }

  explicit BatchFactor(std::vector<FactorType, Allocator>&& factors)
      : factors_(std::move(factors)) {
    updateKeys();
  }

  explicit BatchFactor(const std::vector<FactorType, Allocator>& factors)
      : factors_(factors) {
    updateKeys();
  }

  explicit BatchFactor(const std::vector<FactorType>& factors) {
    factors_.reserve(factors.size());
    factors_.assign(factors.begin(), factors.end());
    updateKeys();
  }

  explicit BatchFactor(std::vector<FactorType>&& factors) {
    factors_.reserve(factors.size());
    for (auto&& f : factors) {
      factors_.push_back(std::move(f));
    }
    updateKeys();
  }

  template <typename Measurement, typename... Args>
  BatchFactor(Key key1, const std::map<Key, Measurement>& measurements,
              const SharedNoiseModel& model, Args&&... args) {
    factors_.reserve(measurements.size());
    for (const auto& [key2, z] : measurements) {
      factors_.push_back(detail::createFactor<FactorType>(
          key1, key2, z, model, std::forward<Args>(args)...));
    }
    updateKeys();
  }

  template <typename Measurement, typename... Args>
  BatchFactor(const std::map<Key, Measurement>& measurements, Key key2,
              const SharedNoiseModel& model, Args&&... args) {
    factors_.reserve(measurements.size());
    for (const auto& [key1, z] : measurements) {
      factors_.push_back(detail::createFactor<FactorType>(
          key1, key2, z, model, std::forward<Args>(args)...));
    }
    updateKeys();
  }


  void print(
      const std::string& s = "",
      const KeyFormatter& keyFormatter = DefaultKeyFormatter) const override {
    Base::print(s, keyFormatter);
    std::cout << "BatchFactor with " << factors_.size()
              << " factors:" << std::endl;
    for (const auto& f : factors_) {
      f.print("", keyFormatter);
    }
  }

  bool equals(const NonlinearFactor& f, double tol = 1e-9) const override {
    const This* p = dynamic_cast<const This*>(&f);
    if (!p || factors_.size() != p->factors_.size()) return false;
    for (size_t i = 0; i < factors_.size(); ++i) {
      if (!factors_[i].equals(p->factors_[i], tol)) return false;
    }
    return true;
  }


  double error(const Values& c) const override {
    double total_error = 0.0;
    for (const auto& f : factors_) {
      total_error += f.error(c);
    }
    return total_error;
  }

  size_t dim() const override { return factors_.size() * ErrorDim; }
 
  void setUseHessianFactor(bool flag) { useHessianFactor_ = flag; }

  std::shared_ptr<GaussianFactor> linearize(
      const Values& values) const override {
    if (factors_.empty()) return std::make_shared<JacobianFactor>();
 
    if (useHessianFactor_) {
      auto jacobian = linearizeToJacobianFactor(values);
      return std::make_shared<HessianFactor>(*jacobian);
    }
 
    constexpr auto factorSlots = std::make_index_sequence<FactorType::N>{};
    if constexpr (hasFixedDimensions(factorSlots)) {
      if (!hasConstrainedNoiseModel_) {
        return linearizeToBatchJacobian(values, factorSlots);
      }
    }
    return linearizeToJacobianFactor(values);
  }

  template <size_t... Is>
  void collectKeys(const FactorType& f, std::index_sequence<Is...>) {
    (keyInfo_.push_back(KeyInfo{
         f.keys()[Is],
         traits<typename FactorType::template ValueType<Is + 1>>::dimension, Is,
         0}),
     ...);
  }

  void updateKeys() {
    constexpr auto factorSlots = std::make_index_sequence<FactorType::N>{};
    allDimensionsFixed_ = hasFixedDimensions(factorSlots);
    hasConstrainedNoiseModel_ = false;
    allNoiseModelsAreUnit_ = true;
    keyDimensions_.clear();
    // 1. Collect all keys and their dimensions
    keyInfo_.clear();
    keyInfo_.reserve(factors_.size() * FactorType::N);
 
    for (const auto& f : factors_) {
      collectKeys(f, std::make_index_sequence<FactorType::N>{});
    }
 
    // 2. Sort and remove duplicates to get unique keys
    std::sort(keyInfo_.begin(), keyInfo_.end(),
              [](const KeyInfo& a, const KeyInfo& b) { return a.key < b.key; });
    auto isDuplicate = [](const KeyInfo& a, const KeyInfo& b) {
      if (a.key != b.key) return false;
      assert(a.dim == b.dim);
      return true;
    };
    auto last = std::unique(keyInfo_.begin(), keyInfo_.end(), isDuplicate);
    keyInfo_.erase(last, keyInfo_.end());
 
    // 3. Fill keys_ and key information
    keys_.clear();
    keys_.reserve(keyInfo_.size());
    DenseIndex offset = 0;
    for (auto& info : keyInfo_) {
      info.offset = offset;
      keys_.push_back(info.key);
      if (info.dim >= 0) {
        offset += static_cast<DenseIndex>(info.dim);
      }
      if (allDimensionsFixed_ && info.dim >= 0) {
        keyDimensions_.push_back(static_cast<size_t>(info.dim));
      }
    }
 
    // Track noise-model traits once at construction time.
    for (const auto& factor : factors_) {
      const auto& model = factor.noiseModel();
      if (!model) continue;
      hasConstrainedNoiseModel_ |= model->isConstrained();
      allNoiseModelsAreUnit_ &= model->isUnit();
    }
 
    // 4. Cache factor indices
    // Since keys_ is sorted, we can use binary search
    indices_.clear();
    indices_.reserve(factors_.size());
    for (const auto& f : factors_) {
      std::array<DenseIndex, FactorType::N> indices_i;
      for (size_t k = 0; k < FactorType::N; ++k) {
        auto it = std::lower_bound(keys_.begin(), keys_.end(), f.keys()[k]);
        indices_i[k] = std::distance(keys_.begin(), it);
      }
      indices_.push_back(indices_i);
    }
  }
};
 
}  // namespace gtsam