Reverted restart to inline in KrylovSchur

Making constructUR to use basisRotate()

Grid/algorithms/iterative/BlockKrylovSchur.h

@@ -97,6 +97,10 @@ protected:
   GridBase*                  Grid_;
   RitzFilter                 ritzFilter;
 
+  // Prefix used in log messages; derived classes override in their constructor
+  // (e.g. "HarmonicBlockKrylovSchur") so the shared code reports the right name.
+  std::string className = "BlockKrylovSchur";
+
   // Flat storage: basis[s*Nblock + t] is the t-th vector of block s
   // After construction: basis has Nm entries
   std::vector<Field> basis;
@@ -146,6 +150,8 @@ public:
       beta_k(0.0), rtol(0.0)
   {}
 
+  virtual ~BlockKrylovSchur() = default;
+
   //--------------------------------------------------------------------
   // Main entry point
   //--------------------------------------------------------------------
@@ -161,7 +167,7 @@ public:
    * _Nstop   : stop after _Nstop eigenvalues converged
    * _Nblock  : block size
    */
-  void operator()(const std::vector<Field>& v0, int _maxIter, int _Nm, int _Nk,
+  virtual void operator()(const std::vector<Field>& v0, int _maxIter, int _Nm, int _Nk,
                           int _Nstop, int _Nblock = 1, bool doubleOrthog = true,
                           bool doVerify = false)
   {
@@ -176,7 +182,7 @@ public:
       if (useParityFlip) divisor *= 2;
       if (useGamma5)     divisor *= 2;
       assert(Nblock % divisor == 0 && (int)v0.size() >= Nblock / divisor);
-       std::cout << GridLogMessage << "BlockKrylovSchur: divisor = " <<divisor << std::endl;
+      std::cout << GridLogMessage << className << ": divisor = " << divisor << std::endl;
     }
     assert(Nm % Nblock == 0);
     assert(Nk % Nblock == 0);
@@ -189,45 +195,18 @@ public:
     // Approximate largest eigenvalue for tolerance normalisation
     RealD approxLambdaMax = approxMaxEval(v0[0]);
     rtol = Tolerance * approxLambdaMax;
-    std::cout << GridLogMessage << "BlockKrylovSchur: approx max eval = "
+    std::cout << GridLogMessage << className << ": approx max eval = "
               << approxLambdaMax << ", rtol = " << rtol << std::endl;
 
     // Initialise
     H = CMat::Zero(N, N);
     B = CMat::Zero(N, Nblock);
 
-    // Build the starting block by expanding each seed with the requested
-    // symmetry partners.  Order: {v, flip(v)} then doubled by {·, g5(·)}
-    // giving {v, flip(v), g5(v), g5(flip(v))} when both flags are set.
-    int divisor = (useParityFlip ? 2 : 1) * (useGamma5 ? 2 : 1);
-    std::cout << GridLogMessage << "BlockKrylovSchur: divisor = " <<divisor << std::endl;
     int start = 0;
-    std::vector<Field> startBlock;
+    std::vector<Field> startBlock = expandStartBlock(v0);
-    startBlock.reserve(Nblock);
-    for (int i = 0; i < Nblock / divisor; i++) {
-      std::vector<Field> group;
-      group.push_back(v0[i]);
-      if (useParityFlip) {
-        int n = (int)group.size();
-        for (int j = 0; j < n; j++) {
-          Field fp(Grid_);
-          parityFlippedField(group[j], fp);
-          group.push_back(std::move(fp));
-        }
-      }
-      if (useGamma5) {
-        int n = (int)group.size();
-        for (int j = 0; j < n; j++) {
-          Field g5v(Grid_);
-          gamma5Func(group[j], g5v);
-          group.push_back(std::move(g5v));
-        }
-      }
-      for (auto& f : group) startBlock.push_back(std::move(f));
-    }
 
     for (int iter = 0; iter < MaxIter; iter++) {
-      std::cout << GridLogMessage << "BlockKrylovSchur: restart iteration " << iter << std::endl;
+      std::cout << GridLogMessage << className << ": restart iteration " << iter << std::endl;
 
       // ---- Block Arnoldi: extend from block start to block Nm/Nblock ----
       blockArnoldiIteration(startBlock, Nm/Nblock, start, doubleOrthog);
@@ -240,15 +219,12 @@ public:
         verify(lbl);
       }
 
-      // ---- Schur decompose H ----
+      // ---- Restart rotation: compute Q and the rotated+reordered Rayleigh
-      ComplexSchurDecomposition schur(H, false, ritzFilter);
+      //      quotient Hnew.  Base class uses the plain Schur decomposition of
-      std::cout << GridLogMessage << "BlockKrylovSchur: Schur decomposed." << std::endl;
+      //      H; derived classes (e.g. HarmonicBlockKrylovSchur) override this
-
+      //      to target a shift.
-      // Reorder: bring wanted Nk Schur values to top-left
+      CMat Q, Hnew;
-      schur.schurReorder(Nk);
+      restartRotation(Q, Hnew);
-      std::cout << GridLogMessage << "BlockKrylovSchur: Schur reordered." << std::endl;
-
-      CMat Q  = schur.getMatrixQ();
       CMat Qt = Q.adjoint();
 
       // Rotate Krylov basis: basis_new[i] = sum_j basis[j] * Qt(j,i)
@@ -256,9 +232,9 @@ public:
       constructUR(basis2, basis, Qt, N);
       basis = basis2;
 
-      // Update b and H
+      // Update B and H
       B = Q * B;
-      H = schur.getMatrixS();
+      H = Hnew;
 
       // ---- Truncate to Nk ----
       int Nkeep = Nk;
@@ -276,7 +252,7 @@ public:
 
       // beta_k = Frobenius norm of the effective coupling
       beta_k = Btmp.norm();
-      std::cout << GridLogMessage << "BlockKrylovSchur: beta_k = " << beta_k << std::endl;
+      std::cout << GridLogMessage << className << ": beta_k = " << beta_k << std::endl;
 
       // Restart: the new starting block is F (the residual block from Arnoldi)
       startBlock = F;
@@ -291,11 +267,11 @@ public:
       computeEigensystem(Hk, Nkeep);
 
       int Nconv = converged(Nkeep);
-      std::cout << GridLogMessage << "BlockKrylovSchur: converged " << Nconv
+      std::cout << GridLogMessage << className << ": converged " << Nconv
                 << " / " << Nstop << std::endl;
 
       if (Nconv >= Nstop || iter == MaxIter - 1) {
-        std::cout << GridLogMessage << "BlockKrylovSchur: done after " << iter
+        std::cout << GridLogMessage << className << ": done after " << iter
                   << " restarts, " << Nconv << " converged." << std::endl;
         std::cout << GridLogMessage << "Eigenvalues: " << evals << std::endl;
 
@@ -306,7 +282,7 @@ public:
             ComplexD eval_est = toStdCmplx(innerProduct(evecs[k], w));
             w -= eval_est * evecs[k];
             RealD res = std::sqrt(norm2(w));
-            std::cout << GridLogMessage << "BlockKrylovSchur: evec[" << k << "]"
+            std::cout << GridLogMessage << className << ": evec[" << k << "]"
                       << "  eval_reported = " << evals[k]
                       << "  eval_est = " << eval_est
                       << "  || A v - eval_est * v || = " << res << std::endl;
@@ -385,7 +361,7 @@ public:
 
     if (nBasis == 0) {
       std::cout << GridLogMessage
-                << "BlockKrylovSchur::verify [" << label
+                << className << "::verify [" << label
                 << "]: basis is empty." << std::endl;
       return;
     }
@@ -394,7 +370,7 @@ public:
     int Nfull = (int)H.rows();
 
     std::cout << GridLogMessage
-              << "======== BlockKrylovSchur::verify [" << label << "] ========" << std::endl;
+              << "======== " << className << "::verify [" << label << "] ========" << std::endl;
     std::cout << GridLogMessage
               << "  nBasis = " << nBasis << "  Nfull = " << Nfull
               << "  Nblock = " << Nblock << "  nF = " << nF << std::endl;
@@ -529,7 +505,69 @@ public:
    */
   virtual void preRun() {}
 
-private:
+  /**
+   * Compute the restart rotation for the current Rayleigh quotient H.
+   *
+   * Returns Q (unitary) and Hnew = Q H Q^dag reordered so that the wanted
+   * Nk eigenvalues occupy the leading block, ready for truncation.
+   *
+   * Base implementation: the plain Schur decomposition of H, reordered by
+   * ritzFilter.  HarmonicBlockKrylovSchur overrides this to Schur-decompose
+   * the shifted quotient (H - sigma I) so that eigenvalues nearest sigma are
+   * selected.
+   */
+  virtual void restartRotation(CMat& Q, CMat& Hnew)
+  {
+    ComplexSchurDecomposition schur(H, false, ritzFilter);
+    std::cout << GridLogMessage << className << ": Schur decomposed." << std::endl;
+    schur.schurReorder(Nk);
+    std::cout << GridLogMessage << className << ": Schur reordered." << std::endl;
+    Q    = schur.getMatrixQ();
+    Hnew = schur.getMatrixS();
+  }
+
+protected:
+
+  //--------------------------------------------------------------------
+  // Starting-block expansion (parity flip / gamma5 partners)
+  //--------------------------------------------------------------------
+  /**
+   * Expand Nblock/divisor seed vectors into the full Nblock starting block by
+   * pairing each seed with its parity-flipped and/or gamma5 partners.  Order:
+   * {v, flip(v)} then doubled by {., g5(.)} giving {v, flip(v), g5(v),
+   * g5(flip(v))} when both flags are set.  Requires Nblock, Grid_ and the
+   * flags to be set.
+   */
+  std::vector<Field> expandStartBlock(const std::vector<Field>& v0)
+  {
+    int divisor = (useParityFlip ? 2 : 1) * (useGamma5 ? 2 : 1);
+    std::vector<Field> startBlock;
+    startBlock.reserve(Nblock);
+    for (int i = 0; i < Nblock / divisor; i++) {
+      std::vector<Field> group;
+      group.push_back(v0[i]);
+      if (useParityFlip) {
+        int n = (int)group.size();
+        for (int j = 0; j < n; j++) {
+          Field fp(Grid_);
+          parityFlippedField(group[j], fp);
+          group.push_back(std::move(fp));
+        }
+      }
+      if (useGamma5) {
+        int n = (int)group.size();
+        for (int j = 0; j < n; j++) {
+          Field g5v(Grid_);
+          gamma5Func(group[j], g5v);
+          group.push_back(std::move(g5v));
+        }
+      }
+      for (auto& f : group) startBlock.push_back(std::move(f));
+    }
+    return startBlock;
+  }
+
+protected:
 
   //--------------------------------------------------------------------
   // Block Arnoldi iteration
@@ -711,7 +749,7 @@ private:
         // deflation: zero this vector
         W[j] = Zero();
         std::cout << GridLogMessage
-                  << "BlockKrylovSchur: deflation at block column " << j
+                  << className << ": deflation at block column " << j
                   << " (norm = " << nrm << ")" << std::endl;
       }
     }
@@ -806,7 +844,7 @@ private:
       CVec Bty  = Bk.adjoint() * yk;                // Nblock-vector
       RealD res = Bty.norm();
       ritzEstimates.push_back(res);
-      std::cout << GridLogMessage << "BlockKrylovSchur: Ritz estimate[" << k
+      std::cout << GridLogMessage << className << ": Ritz estimate[" << k
                 << "] = " << res << "  eval = " << evals[k] << std::endl;
       if (res < rtol) Nconv++;
       else break;

Grid/algorithms/iterative/KrylovSchur.h

@@ -43,6 +43,19 @@ using std::conj;
 inline const ComplexD& toStdCmplx(const ComplexD& c) { return c; }
 #endif
 
+/**
+ * Small dense coefficient matrix of ComplexD (thrust::complex on GPU builds)
+ * for passing to basisRotate, whose device kernel cannot multiply Grid tensors
+ * by std::complex. Must be namespace-scope: nvcc rejects function-local types
+ * as template arguments to functions containing extended device lambdas.
+ */
+struct KSCoeffMat {
+  std::vector<ComplexD> a; int n;
+  KSCoeffMat(int _n) : a(_n*_n), n(_n) {}
+  ComplexD&       operator()(int i, int j)       { return a[i*n+j]; }
+  const ComplexD& operator()(int i, int j) const { return a[i*n+j]; }
+};
+
 /**
  * Options for which Ritz values to keep in implicit restart. TODO move this and utilities into a new file
  */
@@ -392,7 +405,45 @@ class KrylovSchur {
 
         std::cout << GridLogDebug << "b after Arnoldi " << b << std::endl;
 
-        nonHarmonicRestart();
+        // --- Restart: Schur-decompose Rayleigh, reorder by ritzFilter, rotate
+        //     the basis into the Schur vectors, and truncate to the leading Nk.
+        ComplexSchurDecomposition schur (Rayleigh, false, ritzFilter);
+        std::cout << GridLogDebug << "Schur decomp holds? " << schur.checkDecomposition() << std::endl;
+
+        // Rearrange Schur matrix so wanted evals are on top left (like MATLAB's ordschur)
+        std::cout << GridLogMessage << "Reordering Schur eigenvalues" << std::endl;
+        schur.schurReorder(Nk);
+        std::cout << GridLogMessage << "Schur eigenvalues reordered." << std::endl;
+
+        Eigen::MatrixXcd Q = schur.getMatrixQ();
+        Qt = Q.adjoint();                           // TODO should Q be real?
+
+        std::cout << GridLogMessage << "*** ROTATING TO SCHUR BASIS *** " << std::endl;
+
+        // Rotate Krylov basis, b vector, redefine Rayleigh quotient and evecs, and truncate.
+        Rayleigh = schur.getMatrixS();
+        b = Q * b;            // b^\dag = b^\dag * Q^\dag <==> b = Q*b
+
+        std::vector<Field> basis2;
+        constructUR(basis2, basis, Qt, Nm,Nk);
+
+        std::cout << GridLogMessage << "*** TRUNCATING FOR RESTART *** " << std::endl;
+        std::cout << GridLogDebug << "Rayleigh before truncation: " << std::endl << Rayleigh << std::endl;
+
+        Eigen::MatrixXcd RayTmp = Rayleigh(Eigen::seqN(0, Nk), Eigen::seqN(0, Nk));
+        Rayleigh = RayTmp;
+
+        basis = std::vector<Field> (basis2.begin(), basis2.begin() + Nk);
+
+        b = b.head(Nk).eval();
+
+        std::cout << GridLogDebug << "Rayleigh after truncation: " << std::endl << Rayleigh << std::endl;
+
+        checkKSDecomposition();
+
+        // Compute eigensystem of Rayleigh. Note the eigenvectors correspond to the sorted eigenvalues.
+        computeEigensystem(Rayleigh);
+        std::cout << GridLogMessage << "Eigenvalues (first Nk sorted): " << std::endl << evals << std::endl;
 
         if (checkConvergedAndReport(i)) return;
       }
@@ -433,68 +484,9 @@ class KrylovSchur {
 
         std::cout << GridLogDebug << "b after Arnoldi " << b << std::endl;
 
-        harmonicRestart(shiftVal);
+        // --- Harmonic restart: Schur-decompose the shift-augmented Rayleigh
-
+        //     quotient so Ritz values close to shiftVal are reordered to the
-        if (checkConvergedAndReport(i)) return;
+        //     top, then rotate and truncate as in the non-harmonic case.
-      }
-    }
-
-  private:
-
-    /**
-     * Non-harmonic restart step: Schur-decompose Rayleigh, reorder by ritzFilter,
-     * rotate the basis into the Schur vectors, and truncate to the leading Nk.
-     */
-    void nonHarmonicRestart() {
-      // Perform a Schur decomposition on Rayleigh
-      ComplexSchurDecomposition schur (Rayleigh, false, ritzFilter);
-      std::cout << GridLogDebug << "Schur decomp holds? " << schur.checkDecomposition() << std::endl;
-
-      // Rearrange Schur matrix so wanted evals are on top left (like MATLAB's ordschur)
-      std::cout << GridLogMessage << "Reordering Schur eigenvalues" << std::endl;
-      schur.schurReorder(Nk);
-      std::cout << GridLogMessage << "Schur eigenvalues reordered." << std::endl;
-
-      Eigen::MatrixXcd Q = schur.getMatrixQ();
-      Qt = Q.adjoint();                           // TODO should Q be real?
-
-      std::cout << GridLogMessage << "*** ROTATING TO SCHUR BASIS *** " << std::endl;
-
-      // Rotate Krylov basis, b vector, redefine Rayleigh quotient and evecs, and truncate.
-      Rayleigh = schur.getMatrixS();
-      b = Q * b;            // b^\dag = b^\dag * Q^\dag <==> b = Q*b
-
-      std::vector<Field> basis2;
-      constructUR(basis2, basis, Qt, Nm,Nk);
-//      basis = basis2;
-
-      std::cout << GridLogMessage << "*** TRUNCATING FOR RESTART *** " << std::endl;
-      std::cout << GridLogDebug << "Rayleigh before truncation: " << std::endl << Rayleigh << std::endl;
-
-      Eigen::MatrixXcd RayTmp = Rayleigh(Eigen::seqN(0, Nk), Eigen::seqN(0, Nk));
-      Rayleigh = RayTmp;
-
-      std::vector<Field> basisTmp = std::vector<Field> (basis2.begin(), basis2.begin() + Nk);
-      basis = basisTmp;
-
-      Eigen::VectorXcd btmp = b.head(Nk);
-      b = btmp;
-
-      std::cout << GridLogDebug << "Rayleigh after truncation: " << std::endl << Rayleigh << std::endl;
-
-      checkKSDecomposition();
-
-      // Compute eigensystem of Rayleigh. Note the eigenvectors correspond to the sorted eigenvalues.
-      computeEigensystem(Rayleigh);
-      std::cout << GridLogMessage << "Eigenvalues (first Nk sorted): " << std::endl << evals << std::endl;
-    }
-
-    /**
-     * Harmonic restart step: Schur-decompose the shift-augmented Rayleigh quotient
-     * so Ritz values close to `shiftVal` are reordered to the top, then rotate and
-     * truncate exactly as in the non-harmonic case.
-     */
-    void harmonicRestart(RealD shiftVal) {
         Eigen::MatrixXcd temp = Rayleigh;
         for (int m=0;m<Nm;m++) temp(m,m) -= shiftVal;
         Eigen::MatrixXcd RayleighS = temp.inverse(); // (B-tI)^-1
@@ -524,15 +516,12 @@ class KrylovSchur {
         Eigen::MatrixXcd RayTmp_s = Btilde(Eigen::seqN(0, Nk), Eigen::seqN(0, Nk));
         Btilde = RayTmp_s;
 
-      std::vector<Field> basisTmp_s = std::vector<Field> (basis2_s.begin(), basis2_s.begin() + Nk);
+        basis2_s = std::vector<Field> (basis2_s.begin(), basis2_s.begin() + Nk);
-      basis2_s = basisTmp_s;
 
-      Eigen::VectorXcd btmp_s = b_s.head(Nk);
+        b_s = b_s.head(Nk).eval();
-      b_s = btmp_s;
 
         Eigen::VectorXcd ghat = -Q_s * g;
-      Eigen::VectorXcd gtmp_s = ghat.head(Nk);
+        ghat = ghat.head(Nk).eval();
-      ghat = gtmp_s;
 
         Field uhat(Grid);
         uhat  = utilde;
@@ -556,7 +545,12 @@ class KrylovSchur {
         checkKSDecomposition();
         computeEigensystem(Rayleigh);
         std::cout << GridLogMessage << "Eigenvalues (first Nk sorted): " << std::endl << evals << std::endl;
+
+        if (checkConvergedAndReport(i)) return;
       }
+    }
+
+  private:
 
     /**
      * Shared post-restart convergence check. On convergence (or the final
@@ -915,30 +909,22 @@ class KrylovSchur {
     /**
      * Given a vector of fields U (equivalently, a LxN matrix, where L is the number of degrees of
      * freedom on the lattice field) and an NxN matrix R, forms the product UR.
+     * Only the first N2 output columns are computed; columns [N2, N) are Zero.
      *
-     * Note that I believe this is equivalent to basisRotate(U, R.adjoint(), 0, N, 0, N, N), but I'm 
+     * Implemented with the accelerator basisRotate kernel, which computes
-     * not 100% sure (this will be slower and unoptimized though).
+     * basis[j] <- sum_k Qt(j,k) basis[k] in place, so we copy U into UR and
+     * pass R^T. The coefficients are staged as ComplexD (thrust::complex on
+     * GPU builds) since std::complex has no device arithmetic with Grid tensors.
      */
     void constructUR(std::vector<Field>& UR, std::vector<Field> &U, Eigen::MatrixXcd& R, int N, int N2) {
-      Field tmp (Grid);
+      KSCoeffMat Rt(N);
+      for (int i = 0; i < N; i++)
+        for (int j = 0; j < N; j++)
+          Rt(i, j) = ComplexD(R(j, i).real(), R(j, i).imag());
 
-      UR.clear();
+      UR = U;
-      // UR.resize(N);
+      basisRotate(UR, Rt, 0, N2, 0, N, N);
-
+      for (int i = N2; i < N; i++) UR[i] = Zero();
-      std::cout << GridLogDebug << "R to rotate by (should be Rayleigh): " << R << std::endl;
-
-      for (int i = 0; i < N; i++) {
-        tmp = Zero();
-	if (i < N2) 
-        for (int j = 0; j < N; j++) {
-          std::cout << GridLogDebug << "Adding R("<<j<<", "<<i<<") = " << R(j, i) << " to rotated" << std::endl;
-          std::cout << GridLogDebug << "Norm of U[j] is " << norm2(U[j]) << " to rotated" << std::endl;
-          tmp = tmp + U[j] * R(j, i);
-        }
-        std::cout << GridLogDebug << "rotated norm at i = " << i << " is: " << norm2(tmp) << std::endl;
-        UR.push_back(tmp);
-        // UR[i] = tmp;
-      }
       return;
     }