sparse_block_matrix.hpp

namespace NDEVR 
{
    using namespace Eigen;
 
    namespace 
    {
        struct TripletEntry
        {
            int r, c;
            g_type x;
            TripletEntry(int r_, int c_, g_type x_) : r(r_), c(c_), x(x_) {}
        };
        struct TripletColSort
        {
            bool operator()(const TripletEntry& e1, const TripletEntry& e2) const
            {
                return e1.c < e2.c || (e1.c == e2.c && e1.r < e2.r);
            }
        };
        template<class T1, class T2, class Pred = std::less<T1> >
        struct CmpPairFirst {
            bool operator()(const std::pair<T1,T2>& left, const std::pair<T1,T2>& right) {
                return Pred()(left.first, right.first);
            }
        };
    }
 
    template <class MatrixType>
    SparseBlockMatrix<MatrixType>::SparseBlockMatrix( const int * rbi, const int* cbi, int rb, int cb, bool hasStorage)
        : _rowBlockIndices(rbi,rb)
        , _colBlockIndices(cbi,cb)
        , _blockCols(cast<uint04>(cb), IntBlockMap())
        , _hasStorage(hasStorage) 
    {
    }
 
    template <class MatrixType>
    SparseBlockMatrix<MatrixType>::SparseBlockMatrix( ):
        _blockCols(0), _hasStorage(true) 
    {
    }
 
    template <class MatrixType>
    void SparseBlockMatrix<MatrixType>::clear(bool dealloc) 
    {
#    ifdef G2O_OPENMP
//#  pragma omp parallel for default (shared) if (_blockCols.size() > 100)
#    endif
        for (int i=0; i < static_cast<int>(_blockCols.size()); ++i) 
        {
            if (!_hasStorage || !dealloc)
            {
                if (_blockCols[i].size() == 0)
                    continue;
                for (auto it = _blockCols[i].begin(); it != _blockCols[i].end(); ++it)
                {
                    it->second.setZero();
                }
            }
            if (_hasStorage && dealloc)
                _blockCols[i].clear();
        }
        //_blockCols.clear();
    }
 
    template <class MatrixType>
    SparseBlockMatrix<MatrixType>::~SparseBlockMatrix(){
        if (_hasStorage)
            clear(true);
    }
 
    template <class MatrixType>
    typename SparseBlockMatrix<MatrixType>::SparseMatrixBlock* SparseBlockMatrix<MatrixType>::block(int r, int c, bool alloc) 
    {
        auto it = _blockCols[c].find(r);
        SparseBlockMatrix<MatrixType>::SparseMatrixBlock* _block = nullptr;
        if (it==_blockCols[c].end())
        {
            if (!_hasStorage && ! alloc )
                return nullptr;
            else 
            {
                int rb=rowsOfBlock(r);
                int cb=colsOfBlock(c);
                SparseBlockMatrix<MatrixType>::SparseMatrixBlock block(rb, cb);
                block.setZero();
                _blockCols[c].insert(std::make_pair(r, block));
                _block = &_blockCols[c][r];
            }
        } 
        else 
        {
            _block = &_blockCols[c][r];
        }
        return _block;
    }
 
    template <class MatrixType>
    const typename SparseBlockMatrix<MatrixType>::SparseMatrixBlock* SparseBlockMatrix<MatrixType>::block(int r, int c) const {
        typename SparseBlockMatrix<MatrixType>::IntBlockMap::const_iterator it =_blockCols[c].find(r);
        if (it==_blockCols[c].end())
            return nullptr;
        return &(it->second);
    }
 
 
 
 
    template <class MatrixType>
    template <class MatrixTransposedType>
    bool SparseBlockMatrix<MatrixType>::transpose(SparseBlockMatrix<MatrixTransposedType>*& dest) const {
        if (! dest){
            dest=new SparseBlockMatrix<MatrixTransposedType>(&_colBlockIndices[0], &_rowBlockIndices[0], _colBlockIndices.size(), _rowBlockIndices.size());
        } else {
            if (! dest->_hasStorage)
                return false;
            if(_rowBlockIndices.size()!=dest->_colBlockIndices.size())
                return false;
            if (_colBlockIndices.size()!=dest->_rowBlockIndices.size())
                return  false;
            for (uint04 i=0; i<_rowBlockIndices.size(); ++i){
                if(_rowBlockIndices[i]!=dest->_colBlockIndices[i])
                    return false;
            }
            for (uint04 i=0; i<_colBlockIndices.size(); ++i){
                if(_colBlockIndices[i]!=dest->_rowBlockIndices[i])
                    return false;
            }
        }
 
        for (uint04 i=0; i<_blockCols.size(); ++i)
        {
            for (auto it = _blockCols[i].begin(); it!=_blockCols[i].end(); ++it){
                const typename SparseBlockMatrix<MatrixType>::SparseMatrixBlock& s=it->second;
                typename SparseBlockMatrix<MatrixTransposedType>::SparseMatrixBlock* d=dest->block(i,it->first,true);
                *d = s.transpose();
            }
        }
        return true;
    }
 
    template <class MatrixType>
    bool SparseBlockMatrix<MatrixType>::add(SparseBlockMatrix& dest) const 
    {
        if (!dest._hasStorage)
            return false;
        if (_rowBlockIndices.size() != dest._rowBlockIndices.size())
            return false;
        if (_colBlockIndices.size() != dest._colBlockIndices.size())
            return  false;
        for (uint04 i = 0; i < _rowBlockIndices.size(); ++i)
        {
            if (_rowBlockIndices[i] != dest._rowBlockIndices[i])
                return false;
        }
        for (uint04 i = 0; i < _colBlockIndices.size(); ++i)
        {
            if (_colBlockIndices[i] != dest._colBlockIndices[i])
                return false;
        }
        for (uint04 i = 0; i < _blockCols.size(); ++i)
        {
            for (auto it = _blockCols[i].begin(); it != _blockCols[i].end(); ++it)
            {
                const SparseBlockMatrix<MatrixType>::SparseMatrixBlock& s = it->second;
                typename SparseBlockMatrix<MatrixType>::SparseMatrixBlock* d = dest.block(it->first, i, true);
                (*d) += s;
            }
        }
        return true;
    }
 
    template <class MatrixType>
    template < class MatrixResultType, class MatrixFactorType >
    bool SparseBlockMatrix<MatrixType>::multiply(SparseBlockMatrix<MatrixResultType>*& dest, const SparseBlockMatrix<MatrixFactorType> * M) const {
        // sanity check
        if (_colBlockIndices.size()!=M->_rowBlockIndices.size())
            return false;
        for (uint04 i=0; i<_colBlockIndices.size(); ++i){
            if (_colBlockIndices[i]!=M->_rowBlockIndices[i])
                return false;
        }
        if (! dest) {
            dest=new SparseBlockMatrix<MatrixResultType>(&_rowBlockIndices[0],&(M->_colBlockIndices[0]), _rowBlockIndices.size(), M->_colBlockIndices.size() );
        }
        if (! dest->_hasStorage)
            return false;
        for (uint04 i=0; i<M->_blockCols.size(); ++i)
        {
            for (auto it=M->_blockCols[i].begin(); it!=M->_blockCols[i].end(); ++it){
                // look for a non-zero block in a row of column it
                int colM=i;
                const typename SparseBlockMatrix<MatrixFactorType>::SparseMatrixBlock& b=it->second;
                typename SparseBlockMatrix<MatrixType>::IntBlockMap::const_iterator rbt=_blockCols[it->first].begin();
                while(rbt!=_blockCols[it->first].end()){
                    //int colA=it->first;
                    int rowA=rbt->first;
                    const typename SparseBlockMatrix<MatrixType>::SparseMatrixBlock& a=rbt->second;
                    typename SparseBlockMatrix<MatrixResultType>::SparseMatrixBlock *c=dest->block(rowA,colM,true);
                    assert (c->rows()==a.rows());
                    assert (c->cols()==b.cols());
                    ++rbt;
                    (*c)+=a*b;
                }
            }
        }
        return true;
    }
 
    template <class MatrixType>
    void SparseBlockMatrix<MatrixType>::multiply(g_type*& dest, const g_type* src) const {
        if (! dest){
            dest=new g_type [_rowBlockIndices[_rowBlockIndices.size()-1] ];
            memset(dest,0, _rowBlockIndices[_rowBlockIndices.size()-1]*sizeof(g_type));
        }
 
        // map the memory by Eigen
        Eigen::Map<VectorX<g_type>> destVec(dest, rows());
        const Eigen::Map<const VectorX<g_type>> srcVec(src, cols());
 
        for (uint04 i=0; i<_blockCols.size(); ++i)
        {
            int srcOffset = i ? _colBlockIndices[i-1] : 0;
 
            for (auto it=_blockCols[i].begin(); it!=_blockCols[i].end(); ++it){
                const typename SparseBlockMatrix<MatrixType>::SparseMatrixBlock& a=it->second;
                int destOffset = it->first ? _rowBlockIndices[it->first - 1] : 0;
                // destVec += *a * srcVec (according to the sub-vector parts)
                internal::axpy(a, srcVec, srcOffset, destVec, destOffset);
            }
        }
    }
 
    template <class MatrixType>
    void SparseBlockMatrix<MatrixType>::multiplySymmetricUpperTriangle(g_type*& dest, const g_type* src) const
    {
        if (! dest){
            dest=new g_type [_rowBlockIndices[_rowBlockIndices.size()-1] ];
            memset(dest,0, _rowBlockIndices[_rowBlockIndices.size()-1]*sizeof(g_type));
        }
 
        // map the memory by Eigen
        Eigen::Map<VectorX<g_type>> destVec(dest, rows());
        const Eigen::Map<const VectorX<g_type>> srcVec(src, cols());
 
        for (uint04 i=0; i<_blockCols.size(); ++i){
            int srcOffset = colBaseOfBlock(i);
            for (typename SparseBlockMatrix<MatrixType>::IntBlockMap::const_iterator it=_blockCols[i].begin(); it!=_blockCols[i].end(); ++it){
                const auto& a = it->second;
                int destOffset = rowBaseOfBlock(it->first);
                if (destOffset > srcOffset) // only upper triangle
                    break;
                // destVec += *a * srcVec (according to the sub-vector parts)
                internal::axpy(a, srcVec, srcOffset, destVec, destOffset);
                if (destOffset < srcOffset)
                    internal::atxpy(a, srcVec, destOffset, destVec, srcOffset);
            }
        }
    }
 
    template <class MatrixType>
    void SparseBlockMatrix<MatrixType>::rightMultiply(g_type*& dest, const g_type* src) const {
        int destSize=cols();
 
        if (! dest){
            dest=new g_type [ destSize ];
            memset(dest,0, destSize*sizeof(g_type));
        }
 
        // map the memory by Eigen
        Eigen::Map<VectorX<g_type>> destVec(dest, destSize);
        Eigen::Map<const VectorX<g_type>> srcVec(src, rows());
 
#    ifdef G2O_OPENMP
#    pragma omp parallel for default (shared) schedule(dynamic, 10)
#    endif
        for (int i=0; i < static_cast<int>(_blockCols.size()); ++i){
            int destOffset = colBaseOfBlock(i);
            for (typename SparseBlockMatrix<MatrixType>::IntBlockMap::const_iterator it=_blockCols[i].begin(); 
                    it!=_blockCols[i].end(); 
                    ++it){
                const typename SparseBlockMatrix<MatrixType>::SparseMatrixBlock& a=it->second;
                int srcOffset = rowBaseOfBlock(it->first);
                // destVec += *a.transpose() * srcVec (according to the sub-vector parts)
                internal::atxpy(a, srcVec, srcOffset, destVec, destOffset);
            }
        }
        
    }
 
    template <class MatrixType>
    void SparseBlockMatrix<MatrixType>::scale(g_type a_) 
    {
        lib_assert(a_ > -200.0f && a_ < 200.0f, "Bad scale");
        for (uint04 i=0; i<_blockCols.size(); ++i){
            for (auto it=_blockCols[i].begin(); it!=_blockCols[i].end(); ++it){
                it->second *= a_;
            }
        }
    }
 
    template <class MatrixType>
    SparseBlockMatrix<MatrixType>*  SparseBlockMatrix<MatrixType>::slice(int rmin, int rmax, int cmin, int cmax, bool alloc) const {
        int m=rmax-rmin;
        int n=cmax-cmin;
        Buffer<int> rowIdx(m);
        rowIdx[0] = rowsOfBlock(rmin);
        for (int i=1; i<m; ++i){
            rowIdx[i]=rowIdx[i-1]+rowsOfBlock(rmin+i);
        }
 
        Buffer<int> colIdx(n);
        colIdx[0] = colsOfBlock(cmin);
        for (int i=1; i<n; ++i)
        {
            colIdx[i]=colIdx[i-1]+colsOfBlock(cmin+i);
        }
        SparseBlockMatrix* s=new SparseBlockMatrix(rowIdx, colIdx, m, n, true);
        for (int i=0; i<n; ++i){
            int mc=cmin+i;
            for (typename SparseBlockMatrix<MatrixType>::IntBlockMap::const_iterator it=_blockCols[mc].begin(); it!=_blockCols[mc].end(); ++it){
                if (it->first >= rmin && it->first < rmax){
                    s->_blockCols[i].insert(std::make_pair(it->first-rmin, it->second));
                }
            }
        }
        s->_hasStorage=alloc;
        return s;
    }
 
    template <class MatrixType>
    uint04 SparseBlockMatrix<MatrixType>::nonZeroBlocks() const {
        uint04 count=0;
        for (uint04 i=0; i<_blockCols.size(); ++i)
            count+=cast<uint04>(_blockCols[i].size());
        return count;
    }
 
    template <class MatrixType>
    uint04 SparseBlockMatrix<MatrixType>::nonZeros() const
    {
        if constexpr (MatrixType::SizeAtCompileTime != Eigen::Dynamic) 
        {
            uint04 nnz = nonZeroBlocks() * MatrixType::SizeAtCompileTime;
            return nnz;
        } 
        else 
        {
            uint04 count=0;
            for (uint04 i=0; i<_blockCols.size(); ++i)
            {
                for (auto it=_blockCols[i].begin(); it!=_blockCols[i].end(); ++it)
                {
                    const SparseBlockMatrix<MatrixType>::SparseMatrixBlock& a=it->second;
                    count += cast<uint04>(a.cols() * a.rows());
                }
            }
            return count;
        }
    }
 
    template <class MatrixType>
    std::ostream& operator << (std::ostream& os, const SparseBlockMatrix<MatrixType>& m){
        os << "RBI: " << m.rowBlockIndices().size();
        for (uint04 i=0; i<m.rowBlockIndices().size(); ++i)
            os << " " << m.rowBlockIndices()[i];
        os << std::endl;
        os << "CBI: " << m.colBlockIndices().size();
        for (uint04 i=0; i<m.colBlockIndices().size(); ++i)
            os << " " << m.colBlockIndices()[i];
        os << std::endl;
 
        for (uint04 i=0; i<m.blockCols().size(); ++i){
            for (typename SparseBlockMatrix<MatrixType>::IntBlockMap::const_iterator it=m.blockCols()[i].begin(); it!=m.blockCols()[i].end(); ++it){
                const typename SparseBlockMatrix<MatrixType>::SparseMatrixBlock& b=it->second;
                os << "BLOCK: " << it->first << " " << i << std::endl;
                os << b << std::endl;
            }
        }
        return os;
    }
 
    template <class MatrixType>
    bool SparseBlockMatrix<MatrixType>::symmPermutation(SparseBlockMatrix<MatrixType>*& dest, const int* pinv, bool upperTriangle) const{
        // compute the permuted version of the new row/column layout
        uint04 n=_rowBlockIndices.size();
        // computed the block sizes
        Buffer<int> blockSizes;
        blockSizes.resize(_rowBlockIndices.size());
        blockSizes[0]=_rowBlockIndices[0];
        for (uint04 i=1; i<n; ++i){
            blockSizes[i]=_rowBlockIndices[i]-_rowBlockIndices[i-1];
        }
        // permute them
        Buffer<int> pBlockIndices;
        pBlockIndices.resize(_rowBlockIndices.size());
        for (uint04 i=0; i<n; ++i){
            pBlockIndices[pinv[i]]=blockSizes[i];
        }
        for (uint04 i=1; i<n; ++i){
            pBlockIndices[i]+=pBlockIndices[i-1];
        }
        // allocate C, or check the structure;
        if (! dest){
            dest=new SparseBlockMatrix(&pBlockIndices[0], &pBlockIndices[0], n, n);
        } else {
            if (dest->_rowBlockIndices.size()!=n)
                return false;
            if (dest->_colBlockIndices.size()!=n)
                return false;
            for (uint04 i=0; i<n; ++i){
                if (dest->_rowBlockIndices[i]!=pBlockIndices[i])
                    return false;
                if (dest->_colBlockIndices[i]!=pBlockIndices[i])
                    return false;
            }
            dest->clear();
        }
        // now ready to permute the columns
        for (uint04 i=0; i<n; ++i){
            //cerr << PVAR(i) <<    " ";
            int pi=pinv[i];
            for (typename SparseBlockMatrix<MatrixType>::IntBlockMap::const_iterator it=_blockCols[i].begin(); 
                    it!=_blockCols[i].end(); ++it){
                int pj=pinv[it->first];
 
                const typename SparseBlockMatrix<MatrixType>::SparseMatrixBlock& s=it->second;
 
                typename SparseBlockMatrix<MatrixType>::SparseMatrixBlock* b=0;
                if (! upperTriangle || pj<=pi) {
                    b=dest->block(pj,pi,true);
                    assert(b->cols()==s.cols());
                    assert(b->rows()==s.rows());
                    *b=s;
                } else {
                    b=dest->block(pi,pj,true);
                    assert(b);
                    assert(b->rows()==s.cols());
                    assert(b->cols()==s.rows());
                    *b=s.transpose();
                }
            }
            //cerr << endl;
            // within each row, 
        }
        return true;
        
    }
 
    template <class MatrixType>
    int SparseBlockMatrix<MatrixType>::fillCCS(g_type* Cx, bool upperTriangle) const
    {
        assert(Cx && "Target destination is NULL");
        g_type* CxStart = Cx;
        for (uint04 i=0; i<_blockCols.size(); ++i){
            int cstart=i ? _colBlockIndices[i-1] : 0;
            int csize=colsOfBlock(i);
            for (int c=0; c<csize; ++c) {
                for (typename SparseBlockMatrix<MatrixType>::IntBlockMap::const_iterator it=_blockCols[i].begin(); it!=_blockCols[i].end(); ++it){
                    const auto& b=it->second;
                    int rstart=it->first ? _rowBlockIndices[it->first-1] : 0;
 
                    int elemsToCopy = cast<int>(b.rows());
                    if (upperTriangle && rstart == cstart)
                        elemsToCopy = c + 1;
                    memcpy(Cx, b.data() + c * b.rows(), elemsToCopy * sizeof(g_type));
                    Cx += cast<size_t>(elemsToCopy);
 
                }
            }
        }
        return cast<int>(Cx - CxStart);
    }
 
    template <class MatrixType>
    int SparseBlockMatrix<MatrixType>::fillCCS(int* Cp, int* Ci, g_type* Cx, bool upperTriangle) const
    {
        assert(Cp && Ci && Cx && "Target destination is NULL");
        int nz=0;
        for (uint04 i=0; i<_blockCols.size(); ++i){
            int cstart=i ? _colBlockIndices[i-1] : 0;
            int csize=colsOfBlock(i);
            for (int c=0; c<csize; ++c) {
                *Cp=nz;
                for (typename SparseBlockMatrix<MatrixType>::IntBlockMap::const_iterator it=_blockCols[i].begin(); it!=_blockCols[i].end(); ++it){
                    const typename SparseBlockMatrix<MatrixType>::SparseMatrixBlock& b=it->second;
                    int rstart=it->first ? _rowBlockIndices[it->first-1] : 0;
 
                    int elemsToCopy = b.rows();
                    if (upperTriangle && rstart == cstart)
                        elemsToCopy = c + 1;
                    for (int r=0; r<elemsToCopy; ++r){
                        *Cx++ = b(r,c);
                        *Ci++ = rstart++;
                        ++nz;
                    }
                }
                ++Cp;
            }
        }
        *Cp=nz;
        return nz;
    }
 
    template <class MatrixType>
    void SparseBlockMatrix<MatrixType>::fillBlockStructure(MatrixStructure& ms) const
    {
        int n        = _colBlockIndices.size();
        int nzMax = (int)nonZeroBlocks();
 
        ms.alloc(n, nzMax);
        ms.m = _rowBlockIndices.size();
 
        int nz = 0;
        int* Cp = ms.Ap;
        int* Ci = ms.Aii;
        for (int i = 0; i < static_cast<int>(_blockCols.size()); ++i){
            *Cp = nz;
            const int& c = i;
            for (typename SparseBlockMatrix<MatrixType>::IntBlockMap::const_iterator it=_blockCols[i].begin(); it!=_blockCols[i].end(); ++it) {
                const int& r = it->first;
                if (r <= c) {
                    *Ci++ = r;
                    ++nz;
                }
            }
            Cp++;
        }
        *Cp=nz;
        assert(nz <= nzMax);
    }
 
    template <class MatrixType>
    int SparseBlockMatrix<MatrixType>::fillSparseBlockMatrixCCS(SparseBlockMatrixCCS<MatrixType>& blockCCS) const
    {
        blockCCS.columns().resize(blockCols().size());
        uint04 numblocks = 0;
        for (uint04 i = 0; i < blockCols().size(); ++i) 
        {
            const IntBlockMap& row = blockCols()[i];
            auto& dest = blockCCS.columns()[i];
            dest.clear();
            dest.setSize(cast<uint04>(row.size()));
            uint04 index = 0;
            for (auto it = row.begin(); it != row.end(); ++it) 
            {
                dest[index++] = typename SparseBlockMatrixCCS<MatrixType>::RowBlock(it->first, it->second);
            }
            numblocks += cast<uint04>(row.size());
        }
        return numblocks;
    }
 
    template <class MatrixType>
    int SparseBlockMatrix<MatrixType>::fillSparseBlockMatrixCCSTransposed(SparseBlockMatrixCCS<MatrixType>& blockCCS) const
    {
        blockCCS.columns().clear();
        blockCCS.columns().resize(_rowBlockIndices.size());
        int numblocks = 0;
        for (uint04 i = 0; i < blockCols().size(); ++i) 
        {
            const IntBlockMap& row = blockCols()[i];
            for (auto it = row.begin(); it != row.end(); ++it) 
            {
                auto& dest = blockCCS.columns()[it->first];
                dest.add(typename SparseBlockMatrixCCS<MatrixType>::RowBlock(i, it->second));
                ++numblocks;
            }
        }
        return numblocks;
    }
 
    template <class MatrixType>
    void SparseBlockMatrix<MatrixType>::takePatternFromHash(SparseBlockMatrixHashMap<MatrixType>& hashMatrix)
    {
        // sort the sparse columns and add them to the map structures by
        // exploiting that we are inserting a sorted structure
        typedef std::pair<int, MatrixType> SparseColumnPair;
        typedef typename SparseBlockMatrixHashMap<MatrixType>::SparseColumn HashSparseColumn;
        for (uint04 i = 0; i < hashMatrix.columns().size(); ++i) {
            // prepare a temporary vector for sorting
            HashSparseColumn& column = hashMatrix.columns()[i];
            if (column.size() == 0)
                continue;
            Buffer<SparseColumnPair> sparseRowSorted; // temporary structure
            sparseRowSorted.ensureCapacity(column.size());
            for (auto it = column.begin(); it != column.end(); ++it)
                sparseRowSorted.add(*it);
            std::sort(sparseRowSorted.begin(), sparseRowSorted.end(), CmpPairFirst<int, MatrixType>());
            // try to free some memory early
            column.clear();
            // now insert sorted vector to the std::map structure
            IntBlockMap& destColumnMap = blockCols()[i];
            //destColumnMap.insert(sparseRowSorted[0]);
            for (uint04 j = 0; j < sparseRowSorted.size(); ++j) 
            {
                //typename SparseBlockMatrix<MatrixType>::IntBlockMap::iterator hint = destColumnMap.end();
                //--hint; // cppreference says the element goes after the hint (until C++11)
                destColumnMap.insert(sparseRowSorted[j]);
            }
        }
    }
 
}// end namespace