eig_block_trlm.cpp

Go to the documentation of this file.

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <iostream>
#include <vector>
#include <algorithm>
 
#include <quda_internal.h>
#include <eigensolve_quda.h>
#include <qio_field.h>
#include <color_spinor_field.h>
#include <blas_quda.h>
#include <util_quda.h>
#include <eigen_helper.h>
 
namespace quda
{
  // Thick Restarted Block Lanczos Method constructor
  BLKTRLM::BLKTRLM(const DiracMatrix &mat, QudaEigParam *eig_param, TimeProfile &profile) :
    TRLM(mat, eig_param, profile)
  {
    bool profile_running = profile.isRunning(QUDA_PROFILE_INIT);
    if (!profile_running) profile.TPSTART(QUDA_PROFILE_INIT);
 
    // Block Thick restart specific checks
    if (n_kr < n_ev + 6) errorQuda("n_kr=%d must be greater than n_ev+6=%d\n", n_kr, n_ev + 6);
 
    if (!(eig_param->spectrum == QUDA_SPECTRUM_LR_EIG || eig_param->spectrum == QUDA_SPECTRUM_SR_EIG)) {
      errorQuda("Only real spectrum type (LR or SR) can be passed to the TR Lanczos solver");
    }
 
    if (n_kr % block_size != 0) {
      errorQuda("Block size %d is not a factor of the Krylov space size %d", block_size, n_kr);
    }
 
    if (n_ev % block_size != 0) {
      errorQuda("Block size %d is not a factor of the compressed space %d", block_size, n_ev);
    }
 
    if (block_size == 0) { errorQuda("Block size %d passed to block eigensolver", block_size); }
 
    int n_blocks = n_kr / block_size;
    block_data_length = block_size * block_size;
    int arrow_mat_array_size = block_data_length * n_blocks;
    // Tridiagonal/Arrow matrix
    block_alpha = (Complex *)safe_malloc(arrow_mat_array_size * sizeof(Complex));
    block_beta = (Complex *)safe_malloc(arrow_mat_array_size * sizeof(Complex));
    for (int i = 0; i < arrow_mat_array_size; i++) {
      block_alpha[i] = 0.0;
      block_beta[i] = 0.0;
    }
 
    // Temp storage used in blockLanczosStep
    jth_block = (Complex *)safe_malloc(block_data_length * sizeof(Complex));
 
    if (!profile_running) profile.TPSTOP(QUDA_PROFILE_INIT);
  }
 
  void BLKTRLM::operator()(std::vector<ColorSpinorField *> &kSpace, std::vector<Complex> &evals)
  {
    // In case we are deflating an operator, save the tunechache from the inverter
    saveTuneCache();
 
    // Pre-launch checks and preparation
    //---------------------------------------------------------------------------
    if (getVerbosity() >= QUDA_VERBOSE) queryPrec(kSpace[0]->Precision());
    // Check to see if we are loading eigenvectors
    if (strcmp(eig_param->vec_infile, "") != 0) {
      printfQuda("Loading evecs from file name %s\n", eig_param->vec_infile);
      loadFromFile(mat, kSpace, evals);
      return;
    }
 
    // Check for an initial guess. If none present, populate with rands, then
    // orthonormalise
    // DMH: This is an important step. With block solvers, initial guesses
    //      of block sizes N can be subspaces rich in extremal eigenmodes,
    //      N times more rich than non-blocked solvers.
    //      Final paragraph, IV.B https://arxiv.org/pdf/1902.02064.pdf
    prepareInitialGuess(kSpace);
 
    // Increase the size of kSpace passed to the function, will be trimmed to
    // original size before exit.
    prepareKrylovSpace(kSpace, evals);
 
    // Check for Chebyshev maximum estimation
    checkChebyOpMax(mat, kSpace);
 
    // Convergence and locking criteria
    double mat_norm = 0.0;
    double epsilon = setEpsilon(kSpace[0]->Precision());
 
    // Print Eigensolver params
    printEigensolverSetup();
    //---------------------------------------------------------------------------
 
    // Begin BLOCK TRLM Eigensolver computation
    //---------------------------------------------------------------------------
    profile.TPSTART(QUDA_PROFILE_COMPUTE);
 
    // Loop over restart iterations.
    while (restart_iter < max_restarts && !converged) {
 
      for (int step = num_keep; step < n_kr; step += block_size) blockLanczosStep(kSpace, step);
      iter += (n_kr - num_keep);
 
      // Solve current block tridiag
      profile.TPSTOP(QUDA_PROFILE_COMPUTE);
      eigensolveFromBlockArrowMat();
      profile.TPSTART(QUDA_PROFILE_COMPUTE);
 
      // mat_norm is updated.
      for (int i = num_locked; i < n_kr; i++)
        if (fabs(alpha[i]) > mat_norm) mat_norm = fabs(alpha[i]);
 
      // Locking check
      iter_locked = 0;
      for (int i = 1; i < (n_kr - num_locked); i++) {
        if (residua[i + num_locked] < epsilon * mat_norm) {
          if (getVerbosity() >= QUDA_DEBUG_VERBOSE)
            printfQuda("**** Locking %d resid=%+.6e condition=%.6e ****\n", i, residua[i + num_locked],
                       epsilon * mat_norm);
          iter_locked = i;
        } else {
          // Unlikely to find new locked pairs
          break;
        }
      }
 
      // Convergence check
      iter_converged = iter_locked;
      for (int i = iter_locked + 1; i < n_kr - num_locked; i++) {
        if (residua[i + num_locked] < tol * mat_norm) {
          if (getVerbosity() >= QUDA_DEBUG_VERBOSE)
            printfQuda("**** Converged %d resid=%+.6e condition=%.6e ****\n", i, residua[i + num_locked], tol * mat_norm);
          iter_converged = i;
        } else {
          // Unlikely to find new converged pairs
          break;
        }
      }
 
      // In order to maintain the block structure, we truncate the
      // algorithmic variables to be multiples of the block size
      iter_keep = std::min(iter_converged + (n_kr - num_converged) / 2, n_kr - num_locked - 12);
      iter_keep = (iter_keep / block_size) * block_size;
      profile.TPSTOP(QUDA_PROFILE_COMPUTE);
      computeBlockKeptRitz(kSpace);
      profile.TPSTART(QUDA_PROFILE_COMPUTE);
 
      num_converged = num_locked + iter_converged;
      num_keep = num_locked + iter_keep;
      num_locked += iter_locked;
 
      // In order to maintain the block structure, we truncate the
      // algorithmic variables to be multiples of the block size
      num_converged = (num_converged / block_size) * block_size;
      num_keep = (num_keep / block_size) * block_size;
      num_locked = (num_locked / block_size) * block_size;
 
      if (getVerbosity() >= QUDA_VERBOSE) {
        printfQuda("%04d converged eigenvalues at restart iter %04d\n", num_converged, restart_iter + 1);
      }
 
      if (getVerbosity() >= QUDA_DEBUG_VERBOSE) {
        printfQuda("iter Conv = %d\n", iter_converged);
        printfQuda("iter Keep = %d\n", iter_keep);
        printfQuda("iter Lock = %d\n", iter_locked);
        printfQuda("num_converged = %d\n", num_converged);
        printfQuda("num_keep = %d\n", num_keep);
        printfQuda("num_locked = %d\n", num_locked);
        for (int i = 0; i < n_kr; i++) {
          printfQuda("Ritz[%d] = %.16e residual[%d] = %.16e\n", i, alpha[i], i, residua[i]);
        }
      }
 
      // Check for convergence
      if (num_converged >= n_conv) converged = true;
      restart_iter++;
    }
 
    profile.TPSTOP(QUDA_PROFILE_COMPUTE);
 
    // Post computation report
    //---------------------------------------------------------------------------
    if (!converged) {
      if (eig_param->require_convergence) {
        errorQuda("BLOCK TRLM failed to compute the requested %d vectors with a %d search space, %d block size, "
                  "and %d Krylov space in %d restart steps. Exiting.",
                  n_conv, n_ev, n_kr, block_size, max_restarts);
      } else {
        warningQuda("BLOCK TRLM failed to compute the requested %d vectors with a %d search space, %d block size, "
                    "and %d Krylov space in %d restart steps. Continuing with current lanczos factorisation.",
                    n_conv, n_ev, n_kr, block_size, max_restarts);
      }
    } else {
      if (getVerbosity() >= QUDA_SUMMARIZE) {
        printfQuda("BLOCK TRLM computed the requested %d vectors in %d restart steps with %d block size and "
                   "%d BLOCKED OP*x operations.\n",
                   n_conv, restart_iter, block_size, iter / block_size);
 
        // Dump all Ritz values and residua
        for (int i = 0; i < n_conv; i++) {
          printfQuda("RitzValue[%04d]: (%+.16e, %+.16e) residual %.16e\n", i, alpha[i], 0.0, residua[i]);
        }
      }
 
      // Compute eigenvalues
      computeEvals(mat, kSpace, evals);
    }
 
    // Local clean-up
    cleanUpEigensolver(kSpace, evals);
  }
 
  // Destructor
  BLKTRLM::~BLKTRLM()
  {
    host_free(jth_block);
    host_free(block_alpha);
    host_free(block_beta);
  }
 
  // Block Thick Restart Member functions
  //---------------------------------------------------------------------------
  void BLKTRLM::blockLanczosStep(std::vector<ColorSpinorField *> v, int j)
  {
    // Compute r = A * v_j - b_{j-i} * v_{j-1}
 
    // Offset for alpha, beta matrices
    int arrow_offset = j * block_size;
    int idx = 0, idx_conj = 0;
 
    // r = A * v_j
    for (int b = 0; b < block_size; b++) chebyOp(mat, *r[b], *v[j + b]);
 
    // r = r - b_{j-1} * v_{j-1}
    int start = (j > num_keep) ? j - block_size : 0;
    if (j - start > 0) {
 
      std::vector<ColorSpinorField *> r_;
      r_.reserve(block_size);
      for (int i = 0; i < block_size; i++) r_.push_back(r[i]);
 
      int blocks = (j - start) / block_size;
      std::vector<Complex> beta_;
      beta_.reserve(blocks * block_data_length);
 
      // Switch beta block order from COLUMN to ROW major
      // This switches the block from upper to lower triangular
      for (int i = 0; i < blocks; i++) {
        int block_offset = (i + start / block_size) * block_data_length;
        for (int b = 0; b < block_size; b++) {
          for (int c = 0; c < block_size; c++) {
            idx = c * block_size + b;
            beta_.push_back(-block_beta[block_offset + idx]);
          }
        }
      }
 
      std::vector<ColorSpinorField *> v_;
      v_.reserve(j - start);
      for (int i = start; i < j; i++) { v_.push_back(v[i]); }
      if (blocks == 1)
        blas::caxpy_L(beta_.data(), v_, r_);
      else
        blas::caxpy(beta_.data(), v_, r_);
    }
 
    // a_j = v_j^dag * r
    std::vector<ColorSpinorField *> vecs_ptr;
    vecs_ptr.reserve(block_size);
    for (int b = 0; b < block_size; b++) { vecs_ptr.push_back(v[j + b]); }
    // Block dot products stored in alpha_block.
    blas::cDotProduct(block_alpha + arrow_offset, vecs_ptr, r);
 
    // Use jth_block to negate alpha data and apply block BLAS.
    // Data is in square hermitian form, no need to switch to ROW major
    for (int b = 0; b < block_size; b++) {
      for (int c = 0; c < block_size; c++) {
        idx = b * block_size + c;
        jth_block[idx] = -1.0 * block_alpha[arrow_offset + idx];
      }
    }
 
    // r = r - a_j * v_j
    blas::caxpy(jth_block, vecs_ptr, r);
 
    // Orthogonalise R[0:block_size] against the Krylov space V[0:j + block_size]
    for (int k = 0; k < 1; k++) blockOrthogonalize(v, r, j + block_size);
 
    // QR decomposition via modified Gram-Schmidt
    // NB, QR via modified Gram-Schmidt is numerically unstable.
    // We perform the QR iteratively to recover numerical stability.
    //
    // Q_0 * R_0(V)   -> Q_0 * R_0 = V
    // Q_1 * R_1(Q_0) -> Q_1 * R_1 = V * R_0^-1 -> Q_1 * R_1 * R_0 = V
    // ...
    // Q_k * R_k(Q_{k-1}) -> Q_k * R_k * R_{k-1} * ... * R_0 = V
    //
    // Where the Q_k are orthonormal to MP and (R_k * R_{k-1} * ... * R_0)^1
    // is the matrix that maps V -> Q_k.
 
    // Column major order
    bool orthed = false;
    int k = 0, kmax = 3;
    while (!orthed && k < kmax) {
      // Compute R_{k}
      if (getVerbosity() >= QUDA_DEBUG_VERBOSE) printfQuda("Orthing k = %d\n", k);
      for (int b = 0; b < block_size; b++) {
        double norm = sqrt(blas::norm2(*r[b]));
        blas::ax(1.0 / norm, *r[b]);
        jth_block[b * (block_size + 1)] = norm;
        for (int c = b + 1; c < block_size; c++) {
 
          Complex cnorm = blas::cDotProduct(*r[b], *r[c]);
          blas::caxpy(-cnorm, *r[b], *r[c]);
 
          idx = c * block_size + b;
          idx_conj = b * block_size + c;
 
          jth_block[idx] = cnorm;
          jth_block[idx_conj] = 0.0;
        }
      }
      // Accumulate R_{k} products
      updateBlockBeta(k, arrow_offset);
      orthed = orthoCheck(r, block_size);
      k++;
    }
 
    // Prepare next step.
    // v_{j+1} = r
    for (int b = 0; b < block_size; b++) *v[j + block_size + b] = *r[b];
 
    // Save Lanczos step tuning
    saveTuneCache();
  }
 
  void BLKTRLM::updateBlockBeta(int k, int arrow_offset)
  {
    if (k == 0) {
      // Copy over the jth_block matrix to block beta, Beta = R_0
      int idx = 0;
      for (int b = 0; b < block_size; b++) {
        for (int c = 0; c < b + 1; c++) {
          idx = b * block_size + c;
          block_beta[arrow_offset + idx] = jth_block[idx];
        }
      }
    } else {
      // Compute BetaNew_ac = (R_k)_ab * Beta_bc
      // Use Eigen, it's neater
      MatrixXcd betaN = MatrixXcd::Zero(block_size, block_size);
      MatrixXcd beta = MatrixXcd::Zero(block_size, block_size);
      MatrixXcd Rk = MatrixXcd::Zero(block_size, block_size);
      int idx = 0;
 
      // Populate matrices
      for (int b = 0; b < block_size; b++) {
        for (int c = 0; c < b + 1; c++) {
          idx = b * block_size + c;
          beta(c, b) = block_beta[arrow_offset + idx];
          Rk(c, b) = jth_block[idx];
        }
      }
 
      // Multiply using Eigen
      betaN = Rk * beta;
 
      // Copy back to beta array
      for (int b = 0; b < block_size; b++) {
        for (int c = 0; c < b + 1; c++) {
          idx = b * block_size + c;
          block_beta[arrow_offset + idx] = betaN(c, b);
        }
      }
    }
  }
 
  void BLKTRLM::eigensolveFromBlockArrowMat()
  {
    profile.TPSTART(QUDA_PROFILE_EIGEN);
    int dim = n_kr - num_locked;
    if (dim % block_size != 0) errorQuda("dim = %d modulo block_size = %d != 0", dim, block_size);
    int blocks = dim / block_size;
 
    int arrow_pos = num_keep - num_locked;
    if (arrow_pos % block_size != 0) errorQuda("arrow_pos = %d modulo block_size = %d != 0", arrow_pos, block_size);
    int block_arrow_pos = arrow_pos / block_size;
    int num_locked_offset = (num_locked / block_size) * block_data_length;
 
    // Eigen objects
    MatrixXcd T = MatrixXcd::Zero(dim, dim);
    block_ritz_mat.resize(dim * dim);
    int idx = 0;
 
    // Populate the r and eblocks
    for (int i = 0; i < block_arrow_pos; i++) {
      for (int b = 0; b < block_size; b++) {
        // E block
        idx = i * block_size + b;
        T(idx, idx) = alpha[idx + num_locked];
 
        for (int c = 0; c < block_size; c++) {
          // r blocks
          idx = num_locked_offset + b * block_size + c;
          T(arrow_pos + c, i * block_size + b) = block_beta[i * block_data_length + idx];
          T(i * block_size + b, arrow_pos + c) = conj(block_beta[i * block_data_length + idx]);
        }
      }
    }
 
    // Add the alpha blocks
    for (int i = block_arrow_pos; i < blocks; i++) {
      for (int b = 0; b < block_size; b++) {
        for (int c = 0; c < block_size; c++) {
          idx = num_locked_offset + b * block_size + c;
          T(i * block_size + b, i * block_size + c) = block_alpha[i * block_data_length + idx];
        }
      }
    }
 
    // Add the beta blocks
    for (int i = block_arrow_pos; i < blocks - 1; i++) {
      for (int b = 0; b < block_size; b++) {
        for (int c = 0; c < b + 1; c++) {
          idx = num_locked_offset + b * block_size + c;
          // Sub diag
          T((i + 1) * block_size + c, i * block_size + b) = block_beta[i * block_data_length + idx];
          // Super diag
          T(i * block_size + b, (i + 1) * block_size + c) = conj(block_beta[i * block_data_length + idx]);
        }
      }
    }
 
    // Invert the spectrum due to Chebyshev (except the arrow diagonal)
    if (reverse) {
      for (int b = 0; b < dim; b++) {
        for (int c = 0; c < dim; c++) {
          T(c, b) *= -1.0;
          if (restart_iter > 0)
            if (b == c && b < arrow_pos && c < arrow_pos) T(c, b) *= -1.0;
        }
      }
    }
 
    // Eigensolve the arrow matrix
    SelfAdjointEigenSolver<MatrixXcd> eigensolver;
    eigensolver.compute(T);
 
    // Populate the alpha array with eigenvalues
    for (int i = 0; i < dim; i++) alpha[i + num_locked] = eigensolver.eigenvalues()[i];
 
    // Repopulate ritz matrix: COLUMN major
    for (int i = 0; i < dim; i++)
      for (int j = 0; j < dim; j++) block_ritz_mat[dim * i + j] = eigensolver.eigenvectors().col(i)[j];
 
    for (int i = 0; i < blocks; i++) {
      for (int b = 0; b < block_size; b++) {
        idx = b * (block_size + 1);
        residua[i * block_size + b + num_locked] = fabs(block_beta[n_kr * block_size - block_data_length + idx]
                                                        * block_ritz_mat[dim * (i * block_size + b + 1) - 1]);
      }
    }
 
    profile.TPSTOP(QUDA_PROFILE_EIGEN);
  }
 
  void BLKTRLM::computeBlockKeptRitz(std::vector<ColorSpinorField *> &kSpace)
  {
    int offset = n_kr + block_size;
    int dim = n_kr - num_locked;
 
    // Multi-BLAS friendly array to store part of Ritz matrix we want
    Complex *ritz_mat_keep = (Complex *)safe_malloc((dim * iter_keep) * sizeof(Complex));
    for (int j = 0; j < dim; j++) {
      for (int i = 0; i < iter_keep; i++) { ritz_mat_keep[j * iter_keep + i] = block_ritz_mat[i * dim + j]; }
    }
 
    rotateVecsComplex(kSpace, ritz_mat_keep, offset, dim, iter_keep, num_locked, profile);
 
    // Update residual vectors
    for (int i = 0; i < block_size; i++) std::swap(kSpace[num_locked + iter_keep + i], kSpace[n_kr + i]);
 
    // Compute new r blocks
    // Use Eigen, it's neater
    MatrixXcd beta = MatrixXcd::Zero(block_size, block_size);
    MatrixXcd ri = MatrixXcd::Zero(block_size, block_size);
    MatrixXcd ritzi = MatrixXcd::Zero(block_size, block_size);
    int blocks = iter_keep / block_size;
    int idx = 0;
    int beta_offset = n_kr * block_size - block_data_length;
    int num_locked_offset = num_locked * block_size;
 
    for (int b = 0; b < block_size; b++) {
      for (int c = 0; c < b + 1; c++) {
        idx = b * block_size + c;
        beta(c, b) = block_beta[beta_offset + idx];
      }
    }
    for (int i = 0; i < blocks; i++) {
      for (int b = 0; b < block_size; b++) {
        for (int c = 0; c < block_size; c++) {
          idx = i * block_size * dim + b * dim + (dim - block_size) + c;
          ritzi(c, b) = block_ritz_mat[idx];
        }
      }
 
      ri = beta * ritzi;
      for (int b = 0; b < block_size; b++) {
        for (int c = 0; c < block_size; c++) {
          idx = num_locked_offset + b * block_size + c;
          block_beta[i * block_data_length + idx] = ri(c, b);
        }
      }
    }
 
    host_free(ritz_mat_keep);
 
    // Save Krylov rotation tuning
    saveTuneCache();
  }
 
} // namespace quda