eigensolve_quda.h

Go to the documentation of this file.

#pragma once
 
#include <quda.h>
#include <quda_internal.h>
#include <dirac_quda.h>
#include <color_spinor_field.h>
 
namespace quda
{
 
  // Local enum for the LU axpy block type
  enum blockType { PENCIL, LOWER_TRI, UPPER_TRI };
 
  class EigenSolver
  {
    using range = std::pair<int, int>;
 
  protected:
    const DiracMatrix &mat;
    QudaEigParam *eig_param;
    TimeProfile &profile;
 
    // Problem parameters
    //------------------
    int n_ev;         
    int n_kr;         
    int n_conv;       
    int n_ev_deflate; 
    double tol;       
    bool reverse;     
    char spectrum[3]; 
    // Algorithm variables
    //--------------------
    bool converged;
    int restart_iter;
    int max_restarts;
    int check_interval;
    int batched_rotate;
    int block_size;
    int iter;
    int iter_converged;
    int iter_locked;
    int iter_keep;
    int num_converged;
    int num_locked;
    int num_keep;
 
    std::vector<double> residua;
 
    // Device side vector workspace
    std::vector<ColorSpinorField *> r;
    std::vector<ColorSpinorField *> d_vecs_tmp;
 
    ColorSpinorField *tmp1;
    ColorSpinorField *tmp2;
 
    QudaPrecision save_prec;
 
  public:
    EigenSolver(const DiracMatrix &mat, QudaEigParam *eig_param, TimeProfile &profile);
 
    virtual ~EigenSolver();
 
    virtual bool hermitian() = 0;
 
    virtual void operator()(std::vector<ColorSpinorField *> &kSpace, std::vector<Complex> &evals) = 0;
 
    static EigenSolver *create(QudaEigParam *eig_param, const DiracMatrix &mat, TimeProfile &profile);
 
    void prepareInitialGuess(std::vector<ColorSpinorField *> &kSpace);
 
    void checkChebyOpMax(const DiracMatrix &mat, std::vector<ColorSpinorField *> &kSpace);
 
    void prepareKrylovSpace(std::vector<ColorSpinorField *> &kSpace, std::vector<Complex> &evals);
 
    double setEpsilon(const QudaPrecision prec);
 
    void queryPrec(const QudaPrecision prec);
 
    void printEigensolverSetup();
 
    void cleanUpEigensolver(std::vector<ColorSpinorField *> &kSpace, std::vector<Complex> &evals);
 
    void matVec(const DiracMatrix &mat, ColorSpinorField &out, const ColorSpinorField &in);
 
    void chebyOp(const DiracMatrix &mat, ColorSpinorField &out, const ColorSpinorField &in);
 
    double estimateChebyOpMax(const DiracMatrix &mat, ColorSpinorField &out, ColorSpinorField &in);
 
    void blockOrthogonalize(std::vector<ColorSpinorField *> v, std::vector<ColorSpinorField *> &r, int j);
 
    void orthonormalizeMGS(std::vector<ColorSpinorField *> &v, int j);
 
    bool orthoCheck(std::vector<ColorSpinorField *> v, int j);
 
    void rotateVecs(std::vector<ColorSpinorField *> &kSpace, const double *rot_array, const int offset, const int dim,
                    const int keep, const int locked, TimeProfile &profile);
 
    void rotateVecsComplex(std::vector<ColorSpinorField *> &kSpace, const Complex *rot_array, const int offset,
                           const int dim, const int keep, const int locked, TimeProfile &profile);
 
    void permuteVecs(std::vector<ColorSpinorField *> &kSpace, int *mat, int size);
 
    void blockRotate(std::vector<ColorSpinorField *> &kSpace, double *array, int rank, const range &i, const range &j,
                     blockType b_type);
 
    void blockRotateComplex(std::vector<ColorSpinorField *> &kSpace, Complex *array, int rank, const range &i,
                            const range &j, blockType b_type, int offset);
 
    void blockReset(std::vector<ColorSpinorField *> &kSpace, int js, int je, int offset);
 
    void deflate(std::vector<ColorSpinorField *> &sol, const std::vector<ColorSpinorField *> &src,
                 const std::vector<ColorSpinorField *> &evecs, const std::vector<Complex> &evals,
                 bool accumulate = false) const;
 
    void deflate(ColorSpinorField &sol, const ColorSpinorField &src, const std::vector<ColorSpinorField *> &evecs,
                 const std::vector<Complex> &evals, bool accumulate = false)
    {
      // FIXME add support for mixed-precison dot product to avoid this copy
      if (src.Precision() != evecs[0]->Precision() && !tmp1) {
        ColorSpinorParam param(*evecs[0]);
        tmp1 = ColorSpinorField::Create(param);
      }
      ColorSpinorField *src_tmp = src.Precision() != evecs[0]->Precision() ? tmp1 : const_cast<ColorSpinorField *>(&src);
      blas::copy(*src_tmp, src); // no-op if these alias
      std::vector<ColorSpinorField *> src_ {src_tmp};
      std::vector<ColorSpinorField *> sol_ {&sol};
      deflate(sol_, src_, evecs, evals, accumulate);
    }
 
    void deflateSVD(std::vector<ColorSpinorField *> &sol, const std::vector<ColorSpinorField *> &vec,
                    const std::vector<ColorSpinorField *> &evecs, const std::vector<Complex> &evals,
                    bool accumulate = false) const;
 
    void deflateSVD(ColorSpinorField &sol, const ColorSpinorField &src, const std::vector<ColorSpinorField *> &evecs,
                    const std::vector<Complex> &evals, bool accumulate = false)
    {
      // FIXME add support for mixed-precison dot product to avoid this copy
      if (src.Precision() != evecs[0]->Precision() && !tmp1) {
        ColorSpinorParam param(*evecs[0]);
        tmp1 = ColorSpinorField::Create(param);
      }
      ColorSpinorField *src_tmp = src.Precision() != evecs[0]->Precision() ? tmp1 : const_cast<ColorSpinorField *>(&src);
      blas::copy(*src_tmp, src); // no-op if these alias
      std::vector<ColorSpinorField *> src_ {src_tmp};
      std::vector<ColorSpinorField *> sol_ {&sol};
      deflateSVD(sol_, src_, evecs, evals, accumulate);
    }
 
    void computeSVD(const DiracMatrix &mat, std::vector<ColorSpinorField *> &evecs, std::vector<Complex> &evals);
 
    void computeEvals(const DiracMatrix &mat, std::vector<ColorSpinorField *> &evecs, std::vector<Complex> &evals,
                      int size);
 
    void computeEvals(const DiracMatrix &mat, std::vector<ColorSpinorField *> &evecs, std::vector<Complex> &evals)
    {
      computeEvals(mat, evecs, evals, n_conv);
    }
 
    void loadFromFile(const DiracMatrix &mat, std::vector<ColorSpinorField *> &eig_vecs, std::vector<Complex> &evals);
 
    void sortArrays(QudaEigSpectrumType spec_type, int n, std::vector<Complex> &x, std::vector<Complex> &y);
 
    void sortArrays(QudaEigSpectrumType spec_type, int n, std::vector<double> &x, std::vector<Complex> &y);
 
    void sortArrays(QudaEigSpectrumType spec_type, int n, std::vector<Complex> &x, std::vector<double> &y);
 
    void sortArrays(QudaEigSpectrumType spec_type, int n, std::vector<double> &x, std::vector<double> &y);
  };
 
  class TRLM : public EigenSolver
  {
 
  public:
    TRLM(const DiracMatrix &mat, QudaEigParam *eig_param, TimeProfile &profile);
 
    virtual ~TRLM();
 
    virtual bool hermitian() { return true; } 
    // Variable size matrix
    std::vector<double> ritz_mat;
 
    // Tridiagonal/Arrow matrix, fixed size.
    double *alpha;
    double *beta;
 
    void operator()(std::vector<ColorSpinorField *> &kSpace, std::vector<Complex> &evals);
 
    void lanczosStep(std::vector<ColorSpinorField *> v, int j);
 
    void reorder(std::vector<ColorSpinorField *> &kSpace);
 
    void eigensolveFromArrowMat();
 
    void computeKeptRitz(std::vector<ColorSpinorField *> &kSpace);
 
  };
 
  class BLKTRLM : public TRLM
  {
  public:
    BLKTRLM(const DiracMatrix &mat, QudaEigParam *eig_param, TimeProfile &profile);
 
    virtual ~BLKTRLM();
 
    virtual bool hermitian() { return true; } 
    // Variable size matrix
    std::vector<Complex> block_ritz_mat;
 
    Complex *block_alpha;
    Complex *block_beta;
 
    Complex *jth_block;
 
    int block_data_length;
 
    void operator()(std::vector<ColorSpinorField *> &kSpace, std::vector<Complex> &evals);
 
    void blockLanczosStep(std::vector<ColorSpinorField *> v, int j);
 
    void eigensolveFromBlockArrowMat();
 
    void updateBlockBeta(int k, int arrow_offset);
 
    void computeBlockKeptRitz(std::vector<ColorSpinorField *> &kSpace);
  };
 
  class IRAM : public EigenSolver
  {
 
  public:
    Complex **upperHess;
    Complex **Qmat;
    Complex **Rmat;
 
    IRAM(const DiracMatrix &mat, QudaEigParam *eig_param, TimeProfile &profile);
 
    virtual bool hermitian() { return false; } 
    virtual ~IRAM();
 
    void operator()(std::vector<ColorSpinorField *> &kSpace, std::vector<Complex> &evals);
 
    void arnoldiStep(std::vector<ColorSpinorField *> &v, std::vector<ColorSpinorField *> &r, double &beta, int j);
 
    void eigensolveFromUpperHess(std::vector<Complex> &evals, const double beta);
 
    void rotateBasis(std::vector<ColorSpinorField *> &v, int keep);
 
    void qrShifts(const std::vector<Complex> evals, const int num_shifts);
 
    void qrIteration(Complex **Q, Complex **R);
 
    void reorder(std::vector<ColorSpinorField *> &kSpace, std::vector<Complex> &evals,
                 const QudaEigSpectrumType spec_type);
  };
 
  void arpack_solve(std::vector<ColorSpinorField *> &h_evecs, std::vector<Complex> &h_evals, const DiracMatrix &mat,
                    QudaEigParam *eig_param, TimeProfile &profile);
 
} // namespace quda