diff --git a/symforce/opt/dense_linearizer.cc b/symforce/opt/dense_linearizer.cc new file mode 100644 index 000000000..edfd6a7be --- /dev/null +++ b/symforce/opt/dense_linearizer.cc @@ -0,0 +1,281 @@ +/* ---------------------------------------------------------------------------- + * SymForce - Copyright 2022, Skydio, Inc. + * This source code is under the Apache 2.0 license found in the LICENSE file. + * ---------------------------------------------------------------------------- */ + +#include "./dense_linearizer.h" + +#include + +#include "./internal/linearizer_utils.h" + +namespace sym { + +template +DenseLinearizer::DenseLinearizer(const std::string& name, + const std::vector>& factors, + const std::vector& key_order, + const bool include_jacobians) + : name_(name), + factors_{&factors}, + state_index_{}, + is_initialized_{false}, + include_jacobians_{include_jacobians} { + if (key_order.empty()) { + keys_ = ComputeKeysToOptimize(factors); + } else { + keys_ = key_order; + } +} + +template +bool DenseLinearizer::IsInitialized() const { + return is_initialized_; +} + +template +const std::vector& DenseLinearizer::Keys() const { + return keys_; +} + +template +const std::unordered_map& DenseLinearizer::StateIndex() const { + SYM_ASSERT(IsInitialized()); + return state_index_; +} + +template +using LinearizedDenseFactor = typename DenseLinearizer::LinearizedDenseFactor; + +template +void CopyJacobianFactorToCombined(const Eigen::DenseBase& jacobian, + const std::vector& key_offsets, + const int residual_offset, + DenseLinearization& linearization) { + const int residual_dim = jacobian.rows(); + for (const linearization_offsets_t& offsets : key_offsets) { + linearization.jacobian.block(residual_offset, offsets.combined_offset, residual_dim, + offsets.tangent_dim) = + jacobian.block(0, offsets.factor_offset, residual_dim, offsets.tangent_dim); + } +} + +template +void CopyRhsFactorToCombined(const LinearizedDenseFactor& linearized_dense_factor, + const std::vector& key_offsets, + DenseLinearization& linearization) { + for (const linearization_offsets_t& offsets : key_offsets) { + linearization.rhs.segment(offsets.combined_offset, offsets.tangent_dim) += + linearized_dense_factor.rhs.segment(offsets.factor_offset, offsets.tangent_dim); + } +} + +template +void CopyHessianFactorToCombined(const LinearizedDenseFactor& factor_linearization, + const std::vector& key_offsets, + DenseLinearization& linearization) { + // NOTE(brad): Assumes no repeats in key_offsets (and thus no repeats in a factor's + // OptimizedKeys()). If there is a repeat, then the strict upper triangle of a diagonal + // block will be copied. Not the biggest deal, but it's a waste. + + // Loop over key blocks of lower triangle of factor linearization hessian + for (int col_index = 0; col_index < static_cast(key_offsets.size()); col_index++) { + const linearization_offsets_t& col_block = key_offsets[col_index]; + + // Diagonal block is triangular + linearization.hessian_lower + .block(col_block.combined_offset, col_block.combined_offset, col_block.tangent_dim, + col_block.tangent_dim) + .template triangularView() += + factor_linearization.hessian.block(col_block.factor_offset, col_block.factor_offset, + col_block.tangent_dim, col_block.tangent_dim); + + // strict lower triangle + for (int row_index = col_index + 1; row_index < static_cast(key_offsets.size()); + row_index++) { + const linearization_offsets_t& row_block = key_offsets[row_index]; + + // If keys have reverse order in combined hessian, the combined block will be in the + // upper triangle, in which case we must transpose to get the lower triangle version. + if (row_block.combined_offset < col_block.combined_offset) { + linearization.hessian_lower.block(col_block.combined_offset, row_block.combined_offset, + col_block.tangent_dim, row_block.tangent_dim) += + factor_linearization.hessian + .block(row_block.factor_offset, col_block.factor_offset, row_block.tangent_dim, + col_block.tangent_dim) + .transpose(); + } else { + linearization.hessian_lower.block(row_block.combined_offset, col_block.combined_offset, + row_block.tangent_dim, col_block.tangent_dim) += + factor_linearization.hessian.block(row_block.factor_offset, col_block.factor_offset, + row_block.tangent_dim, col_block.tangent_dim); + } + } + } +} + +/** + * Linearizes this->factors_ at values into linearization. Is different than subsequent runs + * because it additionally: + * - Allocates space for the dense linearized factors + * - Figures out which parts of the problem linearization correspond to which parts of each + * factor linearization + * - Allocates memory for each factor linearization of each shape + * - Performs sanity checks that the outputs of each factor are consistent + * - Computes indices into values for the arguments of each factor + */ +template +void DenseLinearizer::InitialLinearization(const Values& values, + DenseLinearization& linearization) { + // Compute state vector index + int32_t offset = 0; + for (const Key& key : keys_) { + auto entry = values.IndexEntryAt(key); + entry.offset = offset; + state_index_[key.GetLcmType()] = entry; + + offset += entry.tangent_dim; + } + + // N is sum of tangent dim of all keys to optimize. Equal to num cols of jacobian + const int32_t N = offset; + + // Allocate final storage for the combined RHS since it is dense and has a known size before + // linearizing factors. Allocate temporary storage for the residual because the combined residual + // dimension is not known yet. + linearization.rhs.resize(N); + linearization.rhs.setZero(); + linearization.hessian_lower.resize(N, N); + linearization.hessian_lower.template triangularView().setZero(); + + // NOTE(brad): Currently we assume all factors are dense factors. The reason for this is that it + // is simpler to only copy dense factors into a dense linearization, and want to get that version + // working first before worrying about sparse factors. + linearized_dense_factors_.reserve(factors_->size()); + typename internal::LinearizedDenseFactorPool::SizeTracker size_tracker; + + LinearizedDenseFactor factor_linearization{}; + + std::vector combined_residual; + std::vector factor_jacobians; + + // Track these to make sure that all combined keys are touched by at least one factor. + std::unordered_set keys_touched_by_factors; + + // Inside this loop + for (const auto& factor : *factors_) { + for (const auto& key : factor.OptimizedKeys()) { + keys_touched_by_factors.insert(key); + } + + // First, we evaluate into a temporary linearized factor + factor_indices_.push_back(values.CreateIndex(factor.AllKeys()).entries); + factor.Linearize(values, factor_linearization, &factor_indices_.back()); + + factor_keyoffsets_.emplace_back(); + const int32_t tangent_dim = [&] { + int32_t factor_tangent_dim; + std::tie(factor_keyoffsets_.back(), factor_tangent_dim) = + internal::FactorOffsets(values, factor.OptimizedKeys(), + state_index_); + return factor_tangent_dim; + }(); + const std::vector& key_offsets = factor_keyoffsets_.back(); + + if (key_offsets.empty()) { + std::vector input_keys; + for (const Key& key : factor.OptimizedKeys()) { + input_keys.push_back(key.GetLcmType()); + } + + spdlog::warn( + "LM<{}>: Optimizing a factor that touches no optimized keys! Optimized input keys for " + "the factor are: {}", + name_, input_keys); + } + + internal::AssertConsistentShapes(tangent_dim, factor_linearization, include_jacobians_); + + // Make sure a temporary of the right dimension is kept for relinearizations + linearized_dense_factors_.AppendFactorSize(factor_linearization.residual.rows(), + factor_linearization.rhs.rows(), size_tracker); + + // Add contributions of residual and jacobian to temporary vectors because we don't yet know + // their total size. + combined_residual.insert( + combined_residual.cend(), factor_linearization.residual.data(), + factor_linearization.residual.data() + factor_linearization.residual.size()); + if (include_jacobians_) { + factor_jacobians.emplace_back(std::move(factor_linearization.jacobian)); + } + + // Add contribution of rhs and hessian (can be done directly because we know the whole size) + CopyRhsFactorToCombined(factor_linearization, key_offsets, linearization /* mut */); + CopyHessianFactorToCombined(factor_linearization, key_offsets, linearization /* mut */); + } + // Now that we know the full problem size, we can construct the residual and jacobian + linearization.residual = + Eigen::Map>(combined_residual.data(), combined_residual.size()); + if (include_jacobians_) { + linearization.jacobian.resize(combined_residual.size(), N); + linearization.jacobian.setZero(); + + int row_offset = 0; + for (int i = 0; i < static_cast(factor_jacobians.size()); i++) { + const auto& jacobian = factor_jacobians[i]; + CopyJacobianFactorToCombined(jacobian, factor_keyoffsets_[i], row_offset, linearization); + row_offset += jacobian.rows(); + } + } + + 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)); + } + } + } + + linearization.SetInitialized(); +} + +template +void DenseLinearizer::Relinearize(const Values& values, + DenseLinearization& linearization) { + if (is_initialized_) { + // Set rhs & hessian_lower to 0 as they will be built additively + linearization.rhs.setZero(); + linearization.hessian_lower.template triangularView().setZero(); + // The parts of linearization.jacobian that aren't being set are assumed to have already + // been set to 0 by InitialLinearization and to have not been mutated since. + + int residual_offset = 0; + for (int i = 0; i < static_cast(factors_->size()); i++) { + (*factors_)[i].Linearize(values, linearized_dense_factors_.at(i), &factor_indices_[i]); + + const LinearizedDenseFactor& factor_linearization = linearized_dense_factors_.at(i); + const std::vector& key_offsets = factor_keyoffsets_[i]; + const int residual_dim = factor_linearization.residual.size(); + + // Copy factor_linearzation values into linearization + linearization.residual.segment(residual_offset, residual_dim) = factor_linearization.residual; + if (include_jacobians_) { + CopyJacobianFactorToCombined(factor_linearization.jacobian, key_offsets, residual_offset, + linearization /* mut */); + } + CopyRhsFactorToCombined(factor_linearization, key_offsets, linearization /* mut */); + CopyHessianFactorToCombined(factor_linearization, key_offsets, linearization /* mut */); + + residual_offset += residual_dim; + } + + } else { + InitialLinearization(values, linearization /* mut */); + } +} + +} // namespace sym + +template class sym::DenseLinearizer; +template class sym::DenseLinearizer; diff --git a/symforce/opt/dense_linearizer.h b/symforce/opt/dense_linearizer.h new file mode 100644 index 000000000..3a85ff0dc --- /dev/null +++ b/symforce/opt/dense_linearizer.h @@ -0,0 +1,112 @@ +/* ---------------------------------------------------------------------------- + * SymForce - Copyright 2022, Skydio, Inc. + * This source code is under the Apache 2.0 license found in the LICENSE file. + * ---------------------------------------------------------------------------- */ + +#pragma once + +#include +#include +#include + +#include + +#include "./factor.h" +#include "./internal/linearized_dense_factor_pool.h" +#include "./linearization.h" + +namespace sym { + +template +class DenseLinearizer { + public: + using LinearizedDenseFactor = typename Factor::LinearizedDenseFactor; + + /** + * Construct a Linearizer from factors and optional keys + * + * Args: + * name: name of linearizer used for debug information + * factors: Only stores a pointer, MUST be in scope for the lifetime of this object! + * 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. + * include_jacobians: Relinearize only allocates and fills out the jacobian if true. + */ + DenseLinearizer(const std::string& name, const std::vector>& factors, + const std::vector& key_order = {}, bool include_jacobians = false); + + /** + * 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& 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& 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& values, DenseLinearization& linearization); + + private: + // The name of this linearizer to be used for printing debug information. + std::string name_; + const std::vector>* factors_; + std::vector keys_; + std::unordered_map state_index_; + internal::LinearizedDenseFactorPool linearized_dense_factors_; + bool is_initialized_; + bool include_jacobians_; + + // The index for each factor in the values. Cached the first time we linearize, to avoid repeated + // unordered_map lookups + std::vector> factor_indices_; + + // One std::vector per factor in factors_ (same order). + // nth linearization_offsets_t of a vector is the offsets into factor and state vectors of the + // nth key to optimize (NOT linearized key, a factor key may be linearized but not optimized) of + // the corresponding factor. + std::vector> factor_keyoffsets_; + + /** + * Evaluates the linearizations of the factors at values into linearization, caching all values + * needed for relinearization along the way. Specifically: + * - Calculates state_index_ + * - Allocates all factor linearizations in lineared_dense_factors_ + * - Calculates factor_indices_ + * - Calculates factor_keyoffsets_ + * Sets is_initialized_ to true. + * + * NOTE(brad): Things might break if you call it twice. + */ + void InitialLinearization(const Values& values, + DenseLinearization& linearization); +}; + +} // namespace sym + +extern template class sym::DenseLinearizer; +extern template class sym::DenseLinearizer; diff --git a/test/symforce_dense_linearizer_test.cc b/test/symforce_dense_linearizer_test.cc new file mode 100644 index 000000000..494f02a43 --- /dev/null +++ b/test/symforce_dense_linearizer_test.cc @@ -0,0 +1,159 @@ +/* ---------------------------------------------------------------------------- + * SymForce - Copyright 2022, Skydio, Inc. + * This source code is under the Apache 2.0 license found in the LICENSE file. + * ---------------------------------------------------------------------------- */ + +#include + +#include +#include + +#include +#include +#include +#include +#include + +TEST_CASE("Correct hessian when factor key order doesn't match problem order", + "[dense-linearizer]") { + using M3 = Eigen::Matrix3d; + + // J: 2 0 0 JtJ: 5 0 1 + // 0 3 1 0 9 3 + // 1 0 1 1 3 2 + // Chosen so that if test fails, it'll be easier to visually inspect the hessian and JtJ to see + // what went wrong. + const M3 J = (M3() << 2, 0, 0, 0, 3, 1, 1, 0, 1).finished(); + const std::vector keys = {'x', 'y', 'z'}; + const std::vector factors = {sym::Factord::Jacobian( + [&](const double a, const double b, const double c, Eigen::Vector3d* const res, + M3* const jac) { + if (res != nullptr) { + *res = J * Eigen::Vector3d(a, b, c); + } + if (jac != nullptr) { + *jac = J; + } + }, + keys)}; + + const sym::Valuesd values = [&] { + sym::Valuesd out; + out.Set(keys[0], 7); + out.Set(keys[1], 11); + out.Set(keys[2], 13); + return out; + }(); + + // Since we swapped the order of the first two keys, the jacobian and hessian should be + // J: 0 2 0 JtJ: 9 0 3 + // 3 0 1 0 5 1 + // 0 1 1 3 1 2 + sym::DenseLinearizer linearizer("linearizer", factors, {keys[1], keys[0], keys[2]}, + true /* include_jacobians */); + sym::DenseLinearization linearization; + linearizer.Relinearize(values, linearization); + + const M3 expected_J = (M3() << 0, 2, 0, 3, 0, 1, 0, 1, 1).finished(); + const M3 expected_H = (M3() << 9, 0, 3, 0, 5, 1, 3, 1, 2).finished(); + + CHECK(expected_J == linearization.jacobian); + CHECK(expected_H == M3(linearization.hessian_lower.template selfadjointView())); +} + +TEST_CASE("Residual, jacobian, hessian, and rhs are all consistent", "[dense-linearizer]") { + // Abbreviate types for readability + using M24 = Eigen::Matrix; + using V2 = Eigen::Vector2d; + using V4 = Eigen::Vector4d; + using V8 = Eigen::Matrix; + using M84 = Eigen::Matrix; + using M66 = Eigen::Matrix; + + std::mt19937 gen(7919); + + // Set up problem + std::vector factors; + const M24 J24 = sym::StorageOps::Random(gen); + factors.push_back(sym::Factord::Jacobian( + [J24](const V2& a, const V2& b, V2* const res, M24* const jac) { + if (res != nullptr) { + *res = J24 * (V4() << a, b).finished(); + } + if (jac != nullptr) { + *jac = J24; + } + }, + {'y', 'x'})); + const M84 J84 = sym::StorageOps::Random(gen); + factors.push_back(sym::Factord::Jacobian( + [J84](const V2& a, const V2& b, V8* const res, M84* const jac) { + if (res != nullptr) { + *res = J84 * (V4() << a, b).finished(); + } + if (jac != nullptr) { + *jac = J84; + } + }, + {'z', 'y'})); + + const sym::Valuesd values = [] { + sym::Valuesd out; + out.Set('x', V2(2, 3)); + out.Set('y', V2(5, 7)); + out.Set('z', V2(9, 11)); + return out; + }(); + + sym::DenseLinearizer linearizer("linearizer", factors, {'x', 'y', 'z'}, + true /* include_jacobians */); + sym::DenseLinearization linearization; + linearizer.Relinearize(values, linearization); + + const Eigen::Matrix res = linearization.residual; + const Eigen::Matrix jac = linearization.jacobian; + const auto hes = M66(linearization.hessian_lower.template selfadjointView()); + const Eigen::Matrix rhs = linearization.rhs; + + // Check that linearization is self consistent + CHECK((hes - jac.transpose() * jac).cwiseAbs().maxCoeff() < 1e-13); + CHECK((rhs - jac.transpose() * res).cwiseAbs().maxCoeff() < 1e-13); + + // Check that the initial linearization matches subsequent linearizations + linearizer.Relinearize(values, linearization); + CHECK((res - linearization.residual).cwiseAbs().maxCoeff() < 1e-15); + CHECK((jac - linearization.jacobian).cwiseAbs().maxCoeff() < 1e-15); + CHECK((hes - M66(linearization.hessian_lower.template selfadjointView())) + .cwiseAbs() + .maxCoeff() < 1e-15); + CHECK((rhs - linearization.rhs).cwiseAbs().maxCoeff() < 1e-15); +} + +TEST_CASE("Jacobian is not allocated if include_jacobians is false", "[dense-linearizer]") { + using M3 = Eigen::Matrix3d; + using V3 = Eigen::Vector3d; + + // Set up a dummy problem. The details don't matter. + const std::vector factors = {sym::Factord::Jacobian( + [&](const V3 a, V3* const res, M3* const jac) { + if (res != nullptr) { + *res = a; + } + if (jac != nullptr) { + *jac = M3::Identity(); + } + }, + {'x'})}; + + const sym::Valuesd values = [&] { + sym::Valuesd out; + out.Set('x', V3(2, 3, 5)); + return out; + }(); + + sym::DenseLinearizer linearizer("lin", factors, {'x'}, false /* include_jacobians */); + sym::DenseLinearization linearization; + linearizer.Relinearize(values, linearization); + + CHECK(linearization.jacobian.size() == 0); +}