base_multi_edge.h

#pragma once
#include <iomanip>
#include <limits>
#include <NDEVR/Matrix.h>
#include <Eigen/StdVector>
 
#include "base_edge.h"
#include "robust_kernel.h"
 
namespace NDEVR
{
    using namespace Eigen;

    template <sint04 t_dims, typename E>
    class BaseMultiEdge : public BaseEdge<t_dims, E>
    {
    public:
        struct HessianHelper {
            Eigen::Map<MatrixX<g_type>> matrix;  
            bool transposed;          
            HessianHelper() 
                : matrix(0, 0, 0)
                , transposed(false) 
            {}
        };
 
    public:
        static const int Dimension = BaseEdge<t_dims, E>::Dimension;    
        typedef typename BaseEdge<t_dims, E>::Measurement Measurement;      
        typedef MatrixX<g_type>::MapType JacobianType;                      
        typedef typename BaseEdge<t_dims, E>::ErrorVector ErrorVector;      
        typedef typename BaseEdge<t_dims, E>::InformationType InformationType;  
        typedef Eigen::Map<MatrixX<g_type>, MatrixX<g_type>::Flags & Eigen::PacketAccessBit
            ? Eigen::Aligned
            : Eigen::Unaligned> HessianBlockType;

        BaseMultiEdge()
            : BaseEdge<t_dims, E>()
        {}

        virtual uint04 vertexCount() const final override { return _vertices.size(); };
        virtual const HyperGraph::HGVertex* vertex(uint04 i) const final override { return _vertices[i]; };
        virtual HyperGraph::HGVertex* vertex(uint04 i) final override { return _vertices[i]; };
        virtual void setVertex(uint04 i, HyperGraph::HGVertex* v) final override
        {
            lib_assert(i < _vertices.size(), "index out of bounds");
            lib_assert(dynamic_cast<OptimizableGraph::OGVertex*>(v) != nullptr, "invalid vertex type");
            _vertices[i] = cast<OptimizableGraph::OGVertex*>(v);
        };
 

        static inline int computeUpperTriangleIndex(int i, int j)
        {
            int elemsUpToCol = ((j - 1) * j) / 2;
            return elemsUpToCol + i;
        }

        void constructQuadraticForm()
        {
            if (BaseEdge<t_dims, E>::robustKernel())
            {
                g_type error = BaseEdge<t_dims, E>::chi2();
                Eigen::Vector3<g_type> rho;
                BaseEdge<t_dims, E>::robustKernel()->robustify(error, rho);
                computeQuadraticForm(BaseEdge<t_dims, E>::robustInformation(rho),(-_information * _error) * rho[1]);
            }
            else
            {
                computeQuadraticForm(_information, -_information * _error);
            }
        }
 
 
        void linearizeOplusAndConstructQuadraticForm(JacobianWorkspace& jacobianWorkspace) final override
        {
            for (uint04 i = 0; i < _vertices.size(); ++i)
            {
                OptimizableGraph::OGVertex* v = _vertices[i];
                lib_assert(jacobianWorkspace.maxNumVertices() > i, "Bad vertex size");
                lib_assert(jacobianWorkspace.maxDimension() >= cast<uint04>(t_dims * v->dimension()), "Bad vertex dimension");
                new (&_jacobianOplus[i]) JacobianType(jacobianWorkspace.workspaceForVertex(i), t_dims, v->dimension());
            }
            linearizeOplus();
            constructQuadraticForm();
        }
 
        virtual void linearizeOplus()
        {
            const g_type delta = g_type(1e-9);
            const g_type scalar = g_type(1.0) / (2 * delta);
            ErrorVector errorBak;
            ErrorVector errorBeforeNumeric = _error;
            g_type add_vi[16];
            for (uint04 i = 0; i < _vertices.size(); ++i) 
            {
                //Xi - estimate the jacobian numerically
                OptimizableGraph::OGVertex* vi = _vertices[i];
 
                if (vi->fixed())
                    continue;
 
                const int vi_dim = vi->dimension();
                assert(vi_dim >= 0);
                lib_assert(vi_dim < 16, "not enough memory");
                std::fill(add_vi, add_vi + vi_dim, g_type(0.0));
                lib_assert(_jacobianOplus[i].rows() == t_dims && _jacobianOplus[i].cols() == vi_dim, "jacobian cache dimension does not match");
                _jacobianOplus[i].resize(Dimension, vi_dim);
                // add small step along the unit vector in each dimension
                for (int d = 0; d < vi_dim; ++d) {
                    vi->push();
                    add_vi[d] = delta;
                    vi->oplus(add_vi);
                    this->computeError();
                    errorBak = _error;
                    vi->pop();
                    vi->push();
                    add_vi[d] = -delta;
                    vi->oplus(add_vi);
                    this->computeError();
                    errorBak -= _error;
                    vi->pop();
                    add_vi[d] = 0.0;
                    //if ((scalar * errorBak).array().isNaN().any())
                        //throw "waht";
                    _jacobianOplus[i].col(d) = scalar * errorBak;
                } // end dimension
            }
            _error = errorBeforeNumeric;
            //if (!_error.array().isFinite().all())
                //throw "error";
        }
 
        void mapHessianMemory(g_type* d, int i, int j, bool rowMajor) final override
        {
            int idx = computeUpperTriangleIndex(i, j);
            assert(idx < (int)_hessian.size());
            OptimizableGraph::OGVertex* vi = _vertices[i];
            OptimizableGraph::OGVertex* vj = _vertices[j];
            HessianHelper& h = _hessian[idx];
            if (rowMajor)
            {
                if (h.matrix.data() != d || h.transposed != rowMajor)
                    new (&h.matrix) HessianBlockType(d, vj->dimension(), vi->dimension());
            }
            else
            {
                if (h.matrix.data() != d || h.transposed != rowMajor)
                    new (&h.matrix) HessianBlockType(d, vi->dimension(), vj->dimension());
            }
            h.transposed = rowMajor;
        }

        void resize(uint04 size)
        {
            lib_assert(false, "check this");
            _vertices.resize(size);
            if (size > 0)
            {
                uint04 n = _vertices.size();
                uint04 maxIdx = (n * (n - 1)) / 2;
                _hessian.setSize(maxIdx, HessianHelper());
                _jacobianOplus.setSize(size, JacobianType(0, 0, 0));
            }
            else
            {
                _hessian.clear();
                _jacobianOplus.clear();
            }
 
        }

        bool allVerticesFixed() const final override
        {
            for (uint04 i = 0; i < _vertices.size(); ++i)
            {
                if (!_vertices[i]->fixed())
                    return false;
            }
            return true;
        }

        void computeQuadraticForm(const InformationType& omega, const ErrorVector& weightedError)
        {
            for (uint04 i = 0; i < _vertices.size(); ++i)
            {
                OptimizableGraph::OGVertex* from = _vertices[i];
                if (from->fixed())
                    continue;
                const MatrixX<g_type>& A = _jacobianOplus[i];
                MatrixX<g_type> AtO = A.transpose() * omega;
                int fromDim = from->dimension();
                lib_assert(fromDim >= 0, "Dimension");
                Eigen::Map<MatrixX<g_type>> fromMap(from->hessianData(), fromDim, fromDim);
                Eigen::Map<VectorX<g_type>> fromB(from->bData(), fromDim);
 
                // ii block in the hessian
                fromMap.noalias() += AtO * A;
                fromB.noalias() += A.transpose() * weightedError;
 
                // compute the off-diagonal blocks ij for all j
                for (uint04 j = i + 1; j < _vertices.size(); ++j)
                {
                    OptimizableGraph::OGVertex* to = _vertices[j];
                    bool jstatus = !(to->fixed());
                    if (jstatus)
                    {
                        const MatrixX<g_type>& B = _jacobianOplus[j];
                        int idx = computeUpperTriangleIndex(i, j);
                        assert(idx < (int)_hessian.size());
                        HessianHelper& hhelper = _hessian[idx];
                        if (hhelper.transposed) // we have to write to the block as transposed
                            hhelper.matrix.noalias() += B.transpose() * AtO.transpose();
                        else
                            hhelper.matrix.noalias() += AtO * B;
                        //if (hhelper.matrix.array().isNaN().any())
                            //throw "waht";
                    }
                }
            }
        }
    protected:
        using BaseEdge<t_dims, E>::_information;
        using BaseEdge<t_dims, E>::_error;
        Buffer<OptimizableGraph::OGVertex*> _vertices;  
 
        Buffer<HessianHelper> _hessian;                 
        Buffer<JacobianType> _jacobianOplus; 
 
    public:
        EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    };
 
}