v01

include/basalt/linearization/block_diagonal.hpp

@@ -0,0 +1,89 @@
+#pragma once
+
+#include <unordered_map>
+
+#include <Eigen/Dense>
+
+#include <basalt/utils/cast_utils.hpp>
+
+namespace basalt {
+
+// TODO: expand IndexedBlocks to small class / struct that also holds info on
+// block size and number of blocks, so we don't have to pass it around all the
+// time and we can directly implement things link adding diagonal and matrix
+// vector products in this sparse block diagonal matrix.
+
+// map of camera index to block; used to represent sparse block diagonal matrix
+template <typename Scalar>
+using IndexedBlocks =
+    std::unordered_map<size_t,
+                       Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>>;
+
+// scale dimensions of JTJ as you would do for jacobian scaling of J beforehand,
+// with diagonal scaling matrix D: For jacobain we would use JD, so for JTJ we
+// use DJTJD.
+template <typename Scalar>
+void scale_jacobians(
+    IndexedBlocks<Scalar>& block_diagonal, size_t num_blocks, size_t block_size,
+    const Eigen::Matrix<Scalar, Eigen::Dynamic, 1>& scaling_vector) {
+  BASALT_ASSERT(num_blocks * block_size ==
+                unsigned_cast(scaling_vector.size()));
+  for (auto& [cam_idx, block] : block_diagonal) {
+    auto D =
+        scaling_vector.segment(block_size * cam_idx, block_size).asDiagonal();
+    block = D * block * D;
+  }
+}
+
+// add diagonal to block-diagonal matrix; missing blocks are assumed to be 0
+template <typename Scalar>
+void add_diagonal(IndexedBlocks<Scalar>& block_diagonal, size_t num_blocks,
+                  size_t block_size,
+                  const Eigen::Matrix<Scalar, Eigen::Dynamic, 1>& diagonal) {
+  BASALT_ASSERT(num_blocks * block_size == unsigned_cast(diagonal.size()));
+  for (size_t idx = 0; idx < num_blocks; ++idx) {
+    auto [it, is_new] = block_diagonal.try_emplace(idx);
+    if (is_new) {
+      it->second = diagonal.segment(block_size * idx, block_size).asDiagonal();
+    } else {
+      it->second += diagonal.segment(block_size * idx, block_size).asDiagonal();
+    }
+  }
+}
+
+// sum up diagonal blocks in hash map for parallel reduction
+template <typename Scalar>
+class BlockDiagonalAccumulator {
+ public:
+  using VecX = Eigen::Matrix<Scalar, Eigen::Dynamic, 1>;
+  using MatX = Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>;
+
+  inline void add(size_t idx, MatX&& block) {
+    auto [it, is_new] = block_diagonal_.try_emplace(idx);
+    if (is_new) {
+      it->second = std::move(block);
+    } else {
+      it->second += block;
+    }
+  }
+
+  inline void add_diag(size_t num_blocks, size_t block_size,
+                       const VecX& diagonal) {
+    add_diagonal(block_diagonal_, num_blocks, block_size, diagonal);
+  }
+
+  inline void join(BlockDiagonalAccumulator& b) {
+    for (auto& [k, v] : b.block_diagonal_) {
+      auto [it, is_new] = block_diagonal_.try_emplace(k);
+      if (is_new) {
+        it->second = std::move(v);
+      } else {
+        it->second += v;
+      }
+    }
+  }
+
+  IndexedBlocks<Scalar> block_diagonal_;
+};
+
+}  // namespace basalt