QR_Lapack.h

// Copyright Edward Rosten 2012

//All rights reserved.
//
//Redistribution and use in source and binary forms, with or without
//modification, are permitted provided that the following conditions
//are met:
//1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//2. Redistributions in binary form must reproduce the above copyright
//   notice, this list of conditions and the following disclaimer in the
//   documentation and/or other materials provided with the distribution.
//
//THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND OTHER CONTRIBUTORS ``AS IS''
//AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
//IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
//ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR OTHER CONTRIBUTORS BE
//LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
//CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
//SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
//INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
//CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
//ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
//POSSIBILITY OF SUCH DAMAGE.

#ifndef TOON_INCLUDE_QR_LAPACK_H
#define TOON_INCLUDE_QR_LAPACK_H


#include <TooN/TooN.h>
#include <TooN/lapack.h>
#include <algorithm>
#include <utility>

namespace TooN{

template<int Rows=Dynamic, int Cols=Rows, class Precision=double>
class QR_Lapack{

    private:
        static const int square_Size = (Rows>=0 && Cols>=0)?(Rows<Cols?Rows:Cols):Dynamic;

    public: 
        template<int R, int C, class P, class B> 
        QR_Lapack(const Matrix<R,C,P,B>& m, bool p=0)
        :copy(m),tau(square_size()), 
         Q(square_size(), square_size()), 
         do_pivoting(p), 
         pivot(Zeros(m.num_cols()))
        {
            //pivot is set to all zeros, which means all columns are free columns
            //and can take part in column pivoting.

            compute();
        }
        
        const Matrix<Rows, Cols, Precision, ColMajor>& get_R()
        {
            return copy;
        }
        
        const Matrix<square_Size, square_Size, Precision, ColMajor>& get_Q()
        {
            return Q;
        }   

        const Vector<Cols, int>& get_P()
        {
            return pivot;
        }

    private:

        void compute()
        {   
            FortranInteger M = copy.num_rows();
            FortranInteger N = copy.num_cols();
            
            FortranInteger LWORK=-1;
            FortranInteger INFO;
            FortranInteger lda = M;

            Precision size;
            
            //Set up the pivot vector
            if(do_pivoting)
                pivot = Zeros;
            else
                for(int i=0; i < pivot.size(); i++)
                    pivot[i] = i+1;

            
            //Compute the working space
            geqp3_(&M, &N, copy.get_data_ptr(), &lda, pivot.get_data_ptr(), tau.get_data_ptr(), &size, &LWORK, &INFO);

            LWORK = (FortranInteger) size;

            Precision* work = new Precision[LWORK];
            
            geqp3_(&M, &N, copy.get_data_ptr(), &lda, pivot.get_data_ptr(), tau.get_data_ptr(), work, &LWORK, &INFO);


            if(INFO < 0)
                std::cerr << "error in QR, INFO was " << INFO << std::endl;

            //The upper "triangle+" of copy is R
            //The lower right and tau contain enough information to reconstruct Q
            
            //LAPACK provides a handy function to do the reconstruction
            Q = copy.template slice<0,0,square_Size, square_Size>(0,0,square_size(), square_size());
            
            FortranInteger K = square_size();
            M=K;
            N=K;
            lda = K;
            orgqr_(&M, &N, &K, Q.get_data_ptr(), &lda, tau.get_data_ptr(), work, &LWORK, &INFO);

            if(INFO < 0)
                std::cerr << "error in QR, INFO was " << INFO << std::endl;

            delete [] work;
            
            //Now zero out the lower triangle
            for(int r=1; r < square_size(); r++)
                for(int c=0; c<r; c++)
                    copy[r][c] = 0;

            //Now fix the pivot matrix.
            //We need to go from FORTRAN to C numbering. 
            for(int i=0; i < pivot.size(); i++)
                pivot[i]--;
        }

        Matrix<Rows, Cols, Precision, ColMajor> copy;
        Vector<square_Size, Precision> tau;
        Matrix<square_Size, square_Size, Precision, ColMajor> Q;
        bool do_pivoting;
        Vector<Cols, FortranInteger> pivot;
        

        int square_size()
        {
            return std::min(copy.num_rows(), copy.num_cols());  
        }
};

}


#endif