TridiagEigen.h

// Copyright (C) 2015 Yixuan Qiu
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef TRIDIAG_EIGEN_H
#define TRIDIAG_EIGEN_H

#include <armadillo>
#include <stdexcept>
#include "LapackWrapperExtra.h"

template <typename Scalar = double>
class TridiagEigen
{
private:
    typedef arma::Mat<Scalar> Matrix;
    typedef arma::Col<Scalar> Vector;

    int n;
    Vector main_diag;     // Main diagonal elements of the matrix
    Vector sub_diag;      // Sub-diagonal elements of the matrix
    Matrix evecs;         // To store eigenvectors

    bool computed;

public:
    TridiagEigen() :
        n(0), computed(false)
    {}

    TridiagEigen(const Matrix &mat) :
        n(mat.n_rows), computed(false)
    {
        compute(mat);
    }

    void compute(const Matrix &mat)
    {
        if(!mat.is_square())
            throw std::invalid_argument("TridiagEigen: matrix must be square");

        n = mat.n_rows;
        main_diag = mat.diag();
        sub_diag = mat.diag(-1);
        evecs.set_size(n, n);

        char compz = 'I';
        int lwork = -1;
        Scalar lwork_opt;

        int liwork = -1;
        int liwork_opt, info;

        // Query of lwork and liwork
        arma::lapack::stedc(&compz, &n, main_diag.memptr(), sub_diag.memptr(),
                            evecs.memptr(), &n, &lwork_opt, &lwork, &liwork_opt, &liwork, &info);

        if(info == 0)
        {
            lwork = (int) lwork_opt;
            liwork = liwork_opt;
        } else {
            lwork = 1 + 4 * n + n * n;
            liwork = 3 + 5 * n;
        }

        Scalar *work = new Scalar[lwork];
        int *iwork = new int[liwork];

        arma::lapack::stedc(&compz, &n, main_diag.memptr(), sub_diag.memptr(),
                            evecs.memptr(), &n, work, &lwork, iwork, &liwork, &info);

        delete [] work;
        delete [] iwork;

        if(info < 0)
            throw std::invalid_argument("Lapack stedc: illegal value");
        if(info > 0)
            throw std::logic_error("Lapack stedc: failed to compute all the eigenvalues");

        computed = true;
    }

    Vector eigenvalues()
    {
        if(!computed)
            throw std::logic_error("TridiagEigen: need to call compute() first");

        // After calling compute(), main_diag will contain the eigenvalues.
        return main_diag;
    }

    Matrix eigenvectors()
    {
        if(!computed)
            throw std::logic_error("TridiagEigen: need to call compute() first");

        return evecs;
    }
};



#endif // TRIDIAG_EIGEN_H