block_solver.h

#pragma once
#include <Eigen/Core>
#include "solver.h"
#include "linear_solver.h"
#include "sparse_block_matrix.h"
#include "sparse_block_matrix_diagonal.h"
#include "openmp_mutex.h"
 
 
 
namespace NDEVR 
{
    struct IndexScratch
    {
        Buffer<int> block_pose_indices;         
        Buffer<int> block_landmark_indices;     
    public:
        EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    };
    using namespace Eigen;

    template <class t_solver, int _PoseDim, int _LandmarkDim>
    struct BlockSolverTraits
    {
        static const int PoseDim = _PoseDim;
        static const int LandmarkDim = _LandmarkDim;
        typedef Eigen::Matrix<g_type, PoseDim, PoseDim> PoseMatrixType;
        typedef Eigen::Matrix<g_type, LandmarkDim, LandmarkDim> LandmarkMatrixType;
        typedef Eigen::Matrix<g_type, PoseDim, LandmarkDim> PoseLandmarkMatrixType;
        typedef Eigen::Matrix<g_type, PoseDim, 1> PoseVectorType;
        typedef Eigen::Matrix<g_type, LandmarkDim, 1> LandmarkVectorType;
 
        typedef SparseBlockMatrix<PoseMatrixType> PoseHessianType;
        typedef SparseBlockMatrix<LandmarkMatrixType> LandmarkHessianType;
        typedef SparseBlockMatrix<PoseLandmarkMatrixType> PoseLandmarkHessianType;
        typedef t_solver LinearSolverType;
    };

    template <class t_type>
    struct BlockSolverTraits<t_type, Eigen::Dynamic, Eigen::Dynamic>
    {
        static const int PoseDim = Eigen::Dynamic;
        static const int LandmarkDim = Eigen::Dynamic;
        typedef MatrixX<g_type> PoseMatrixType;
        typedef MatrixX<g_type> LandmarkMatrixType;
        typedef MatrixX<g_type> PoseLandmarkMatrixType;
        typedef VectorX<g_type> PoseVectorType;
        typedef VectorX<g_type> LandmarkVectorType;
 
        typedef SparseBlockMatrix<PoseMatrixType> PoseHessianType;
        typedef SparseBlockMatrix<LandmarkMatrixType> LandmarkHessianType;
        typedef SparseBlockMatrix<PoseLandmarkMatrixType> PoseLandmarkHessianType;
        typedef t_type LinearSolverType;
    };
    template <typename Traits>
    class BlockSolver : public Solver
    {
    public:
 
        static const int PoseDim = Traits::PoseDim;
        static const int LandmarkDim = Traits::LandmarkDim;
        typedef typename Traits::PoseMatrixType PoseMatrixType;
        typedef typename Traits::LandmarkMatrixType LandmarkMatrixType;
        typedef typename Traits::PoseLandmarkMatrixType PoseLandmarkMatrixType;
        typedef typename Traits::PoseVectorType PoseVectorType;
        typedef typename Traits::LandmarkVectorType LandmarkVectorType;
 
        typedef typename Traits::PoseHessianType PoseHessianType;
        typedef typename Traits::LandmarkHessianType LandmarkHessianType;
        typedef typename Traits::PoseLandmarkHessianType PoseLandmarkHessianType;
        typedef typename Traits::LinearSolverType LinearSolverType;
 
    public:
        LinearSolverType solver;    
        BlockSolver();
        void clear()
        {
            _numPoses = 0;
            _numLandmarks = 0;
            _sizePoses = 0;
            _sizeLandmarks = 0;
            _doSchur = true;
            _x.clear();
            _b.clear();
        }

        SparseOptimizer& optimizer()
        {
            return solver.optimizer;
        }

        const SparseOptimizer& optimizer() const
        {
            return solver.optimizer;
        }

        bool init(bool online = false);
        bool buildStructure(bool zeroBlocks, IndexScratch& scratch);
        bool updateStructure(const Buffer<HyperGraph::HGVertex*>& vset, const Set<HyperGraph::HGVertex*>& edges);
        bool buildSystem()
        {
            //# pragma omp parallel for default (shared) if (solver.optimizer.indexMapping().size() > 1000)
            for (OptimizableGraph::OGVertex* v : solver.optimizer.indexMapping())
                v->clearQuadraticForm();
            _Hpp.clear();
            if (_doSchur)
            {
                _Hll.clear();
                _Hpl.clear();
            }
 
            // resetting the terms for the pairwise constraints
            // built up the current system by storing the Hessian blocks in the edges and vertices
            uint04 size = solver.optimizer.activeEdges().size();
            // no threading, we do not need to copy the workspace
            JacobianWorkspace& jacobianWorkspace = solver.optimizer.jacobianWorkspace();
            for (uint04 i = 0; i < size; ++i)
            {
                OptimizableGraph::OGEdge* e = solver.optimizer.activeEdges()[i];
                e->linearizeOplusAndConstructQuadraticForm(jacobianWorkspace); // jacobian of the nodes' oplus (manifold)
            }
 
            // flush the current system in a sparse block matrix
            for (uint04 i = 0; i < solver.optimizer.indexMapping().size(); ++i) 
            {
                OptimizableGraph::OGVertex* v = solver.optimizer.indexMapping()[i];
                uint04 iBase = v->colInHessian();
                if (v->marginalized())
                    iBase += _sizePoses;
                lib_assert(_b.size() > iBase, "Bad size");
                v->copyB(_b.begin(iBase));
            }
            return true;
        }

        bool solve()
        {
            if (!_doSchur)
            {
                bool ok = solver.solve(_Hpp, _x, _b);
                return ok;
            }
            _Hschur.clear();
            _Hpp.add(_Hschur);
            _coefficients.setAllToValue(0.0f);
# ifdef G2O_OPENMP
            //# pragma omp parallel for default (shared) schedule(dynamic, 10)
# endif
            for (uint04 idx = 0; idx < _Hll.blockCols().size(); ++idx)
            {
                auto& marginalizeColumn = _Hll.blockCols()[idx];
 
                assert(marginalizeColumn.size() == 1 && "more than one block in _Hll column");
 
                // calculate inverse block for the landmark
                const LandmarkMatrixType& D = marginalizeColumn.begin()->second;
                //if (!D.array().isFinite().all())
                //  throw "bad ex";
                assert(D.rows() == D.cols() && "Error in landmark matrix");
                LandmarkMatrixType& Dinv = _DInvSchur.diagonal()[idx];
                Dinv = D.inverse();
                //if (!Dinv.array().isFinite().all())
                    //throw "ex";
                LandmarkVectorType db(D.rows());
                for (uint04 j = 0; j < D.rows(); ++j)
                    db[j] = _b[_Hll.rowBaseOfBlock(idx) + _sizePoses + j];
                db = Dinv * db;
 
                const auto& landmarkColumn = _HplCCS.columns()[idx];
 
                for (auto it_outer = landmarkColumn.begin(); it_outer != landmarkColumn.end(); ++it_outer)
                {
                    int i1 = it_outer->row;
 
                    const PoseLandmarkMatrixType& Bi = it_outer->block;
 
                    PoseLandmarkMatrixType BDinv = Bi * Dinv;
                    assert(_HplCCS.rowBaseOfBlock(i1) < cast<int>(_sizePoses) && "Index out of bounds");
                    typename PoseVectorType::MapType Bb(&_coefficients[_HplCCS.rowBaseOfBlock(i1)], Bi.rows());
#       ifdef G2O_OPENMP
                    //ScopedOpenMPMutex mutexLock(&_coefficientsMutex[i1]);
#       endif
                    Bb.noalias() += Bi * db;
                    //if (!Bb.array().isFinite().any())
                    //  throw "ex";
                    assert(i1 >= 0 && i1 < static_cast<int>(_HschurTransposedCCS.columns().size()) && "Index out of bounds");
                    auto targetColumnIt = _HschurTransposedCCS.columns()[i1].begin();
 
                    typename SparseBlockMatrixCCS<PoseLandmarkMatrixType>::RowBlock aux(i1);
                    auto it_inner = std::lower_bound(landmarkColumn.begin(), landmarkColumn.end(), aux);
                    for (; it_inner != landmarkColumn.end(); ++it_inner)
                    {
                        int i2 = it_inner->row;
                        const PoseLandmarkMatrixType& Bj = it_inner->block;
                        while (targetColumnIt->row < i2 /*&& targetColumnIt != _HschurTransposedCCS.blockCols()[i1].end()*/)
                            ++targetColumnIt;
                        assert(targetColumnIt != _HschurTransposedCCS.columns()[i1].end() && targetColumnIt->row == i2 && "invalid iterator, something wrong with the matrix structure");
                        PoseMatrixType& Hi1i2 = targetColumnIt->block;//_Hschur.block(i1,i2);
                        Hi1i2.noalias() -= BDinv * Bj.transpose();
                    }
                }
            }
            //cerr << "Solve [marginalize] = " <<   get_monotonic_time()-t << endl;
 
            // _bschur = _b for calling solver, and not touching _b
            memcpy(_bschur, _b, _sizePoses * sizeof(g_type));
            //for (uint04 i = 0; i < _sizePoses; ++i)
                //if (IsInvalid(_coefficients[i]))
                    //throw "ex";
            for (uint04 i = 0; i < _sizePoses; ++i)
                _bschur[i] -= _coefficients[i];
 
 
            bool solvedPoses = solver.solve(_Hschur, _x, _bschur);
 
            //cerr << "Solve [decompose and solve] = " <<   get_monotonic_time()-t << endl;
            if (!solvedPoses)
                return false;
 
            // _x contains the solution for the poses, now applying it to the landmarks to get the new part of the
            // solution;
            g_type* xp = _x.begin();
            g_type* cp = _coefficients;
            lib_assert(_x.size() > _sizePoses, "Bad size");
            g_type* xl = _x.begin(_sizePoses);
            g_type* cl = _coefficients.begin() + _sizePoses;
            g_type* bl = _b.begin(_sizePoses);
 
            // cp = -xp
            for (uint04 i = 0; i < _sizePoses; ++i)
                cp[i] = -xp[i];
 
            // cl = bl
            memcpy(cl, bl, _sizeLandmarks * sizeof(g_type));
 
            // cl = bl - Bt * xp
            //Bt.multiply(cl, cp);
            _HplCCS.rightMultiply(cl, cp);
 
            // xl = Dinv * cl
            memset(xl, 0, _sizeLandmarks * sizeof(g_type));
            _DInvSchur.multiply(xl, cl);
            //_DInvSchur.rightMultiply(xl,cl);
            //cerr << "Solve [landmark delta] = " <<    get_monotonic_time()-t << endl;
 
            return true;
        }
        //bool computeMarginals(SparseBlockMatrix<MatrixXd>& spinv, const Buffer<std::pair<int, int> >& blockIndices);
        bool setLambda(g_type lambda, bool backup = false);
        void restoreDiagonal();
        bool supportsSchur() { return true; }
        bool schur() { return _doSchur; }
        void setSchur(bool s) { _doSchur = s; }
        void multiplyHessian(g_type* dest, const g_type* src) const { _Hpp.multiplySymmetricUpperTriangle(dest, src); }
 
    protected:
        void resize(IndexScratch& scratch, int totalDim);
 
        SparseBlockMatrix<PoseMatrixType> _Hpp;                             
        SparseBlockMatrix<LandmarkMatrixType> _Hll;                         
        SparseBlockMatrix<PoseLandmarkMatrixType> _Hpl;                     
 
        SparseBlockMatrix<PoseMatrixType> _Hschur;                          
        SparseBlockMatrixDiagonal<LandmarkMatrixType> _DInvSchur;           
 
        SparseBlockMatrixCCS<PoseLandmarkMatrixType> _HplCCS;               
        SparseBlockMatrixCCS<PoseMatrixType> _HschurTransposedCCS;          
 
        Buffer<PoseVectorType> _diagonalBackupPose;                         
        Buffer<LandmarkVectorType> _diagonalBackupLandmark;                 
 
        bool _doSchur;                                                      
 
        Buffer<g_type> _coefficients;                                       
        Buffer<g_type> _bschur;                                             
 
        uint04 _numPoses, _numLandmarks;                                    
        uint04 _sizePoses, _sizeLandmarks;                                  
    public:
        EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    };

    template<class t_type>
    class BlockSolve
    {
    public:
        typedef BlockSolver<BlockSolverTraits<t_type, Eigen::Dynamic, Eigen::Dynamic>> BlockSolverX;    
        typedef BlockSolver<BlockSolverTraits<t_type, 6, 3>> BlockSolver_6_3;   
        typedef BlockSolver<BlockSolverTraits<t_type, 7, 3>> BlockSolver_7_3;   
        typedef BlockSolver<BlockSolverTraits<t_type, 3, 2>> BlockSolver_3_2;   
    };
 
} // end namespace
 
#include "block_solver.hpp"