|
| 1 | +#include <gtest/gtest.h> |
| 2 | +#include <complex> |
| 3 | +#include <vector> |
| 4 | +#include <cmath> |
| 5 | +#include <iostream> |
| 6 | +#include <random> |
| 7 | +#include "source_hsolver/diago_lobpcg.h" |
| 8 | + |
| 9 | +// Define complex double type |
| 10 | +using Complex = std::complex<double>; |
| 11 | + |
| 12 | +// Declare LAPACK zheev helper |
| 13 | +extern "C" { |
| 14 | + void zheev_(const char* jobz, const char* uplo, const int* n, Complex* a, const int* lda, double* w, Complex* work, const int* lwork, double* rwork, int* info); |
| 15 | + void zgemm_(const char* transa, const char* transb, const int* m, const int* n, const int* k, |
| 16 | + const Complex* alpha, const Complex* a, const int* lda, |
| 17 | + const Complex* b, const int* ldb, |
| 18 | + const Complex* beta, Complex* c, const int* ldc); |
| 19 | +} |
| 20 | + |
| 21 | +class DiagoLobpcgTest : public testing::Test { |
| 22 | +protected: |
| 23 | + std::vector<Complex> matrix; |
| 24 | + |
| 25 | + // Generate matrix |
| 26 | + // type 0: Deterministic (Original) |
| 27 | + // type 1: Random Diagonally Dominant Complex Hermitian |
| 28 | + void GenerateMatrix(int n, int type) { |
| 29 | + matrix.resize(n * n); |
| 30 | + if (type == 0) { |
| 31 | + for (int j = 0; j < n; ++j) { |
| 32 | + for (int i = 0; i < n; ++i) { |
| 33 | + if (i == j) { |
| 34 | + matrix[j * n + i] = static_cast<double>(i + 1); // Diagonal 1..n |
| 35 | + } else { |
| 36 | + // Off-diagonal |
| 37 | + double val = 1.0 / (std::abs(i - j) + 1.0); |
| 38 | + matrix[j * n + i] = val * 0.1; |
| 39 | + } |
| 40 | + } |
| 41 | + } |
| 42 | + } else if (type == 1) { |
| 43 | + // Random Hermitian matrix |
| 44 | + // Use specific seed for reproducibility |
| 45 | + std::mt19937 gen(42); |
| 46 | + // Diagonal elements: spaced out to ensure good conditioning for basic LOBPCG |
| 47 | + // Off-diagonal: small random values |
| 48 | + std::uniform_real_distribution<> val_dist(-0.5, 0.5); |
| 49 | + |
| 50 | + for (int j = 0; j < n; ++j) { |
| 51 | + // Diagonal (real) |
| 52 | + // j+1 plus small noise. Keeps eigenvalues well separated approx 1.0 apart. |
| 53 | + matrix[j * n + j] = static_cast<double>(j + 1) + val_dist(gen) * 0.5; |
| 54 | + |
| 55 | + // Off-diagonal (complex) |
| 56 | + for (int i = j + 1; i < n; ++i) { |
| 57 | + Complex val(val_dist(gen) * 0.05, val_dist(gen) * 0.05); |
| 58 | + // A(i, j) at matrix[j*n + i] |
| 59 | + matrix[j * n + i] = val; |
| 60 | + // A(j, i) at matrix[i*n + j] |
| 61 | + matrix[i * n + j] = std::conj(val); |
| 62 | + } |
| 63 | + } |
| 64 | + } |
| 65 | + } |
| 66 | + |
| 67 | + void VerifyLobpcg(int n, int nband, double check_tol = 1e-3, double cg_tol = 1e-5) { |
| 68 | + // --------------------------------------------------------- |
| 69 | + // 1. Solve with LAPACK (Gold Standard) |
| 70 | + // --------------------------------------------------------- |
| 71 | + std::vector<Complex> mat_lapack = matrix; // Deep copy |
| 72 | + std::vector<double> ev_lapack(n); |
| 73 | + |
| 74 | + Complex work_query; |
| 75 | + std::vector<double> rwork(3 * n - 2); |
| 76 | + int lwork_query = -1; |
| 77 | + int info = 0; |
| 78 | + int n_val = n; |
| 79 | + |
| 80 | + char jobz = 'N'; |
| 81 | + char uplo = 'U'; |
| 82 | + |
| 83 | + // Query workspace |
| 84 | + zheev_(&jobz, &uplo, &n_val, mat_lapack.data(), &n_val, ev_lapack.data(), &work_query, &lwork_query, rwork.data(), &info); |
| 85 | + |
| 86 | + int lwork = static_cast<int>(work_query.real()) + 1; |
| 87 | + std::vector<Complex> work(lwork); |
| 88 | + |
| 89 | + // Compute |
| 90 | + zheev_(&jobz, &uplo, &n_val, mat_lapack.data(), &n_val, ev_lapack.data(), work.data(), &lwork, rwork.data(), &info); |
| 91 | + |
| 92 | + ASSERT_EQ(info, 0) << "LAPACK zheev computation failed with info=" << info; |
| 93 | + |
| 94 | + // Output LAPACK eigenvalues for debug |
| 95 | + // std::cout << "LAPACK computed eigenvalues (first 5): "; |
| 96 | + // for(int i=0; i<5 && i<n; ++i) std::cout << ev_lapack[i] << " "; |
| 97 | + // std::cout << std::endl; |
| 98 | + |
| 99 | + // --------------------------------------------------------- |
| 100 | + // 2. Solve with LOBPCG |
| 101 | + // --------------------------------------------------------- |
| 102 | + std::vector<double> precondition(n, 1.0); // Identity Preconditioner |
| 103 | + |
| 104 | + int n_max = nband + 5; |
| 105 | + hsolver::DiagoLOBPCG<Complex> lobpcg(precondition.data(), nband, n, n_max); |
| 106 | + |
| 107 | + std::vector<double> ev_lobpcg(nband); |
| 108 | + std::vector<Complex> psi(n * nband); |
| 109 | + |
| 110 | + // Initialize psi with values |
| 111 | + for(auto &val : psi) val = static_cast<double>(rand()) / RAND_MAX; |
| 112 | + |
| 113 | + auto hpsi_func = [&](Complex* in, Complex* out, const int ld, const int nvec) { |
| 114 | + char transa = 'N'; |
| 115 | + char transb = 'N'; |
| 116 | + int m_ = n; |
| 117 | + int n_ = nvec; |
| 118 | + int k_ = n; |
| 119 | + Complex alpha = 1.0; |
| 120 | + Complex beta = 0.0; |
| 121 | + int lda = n; |
| 122 | + |
| 123 | + zgemm_(&transa, &transb, &m_, &n_, &k_, |
| 124 | + &alpha, matrix.data(), &lda, |
| 125 | + in, &ld, |
| 126 | + &beta, out, &ld); |
| 127 | + }; |
| 128 | + |
| 129 | + int max_iter = 2000; |
| 130 | + bool converged = lobpcg.diag( |
| 131 | + hpsi_func, |
| 132 | + nullptr, |
| 133 | + false, |
| 134 | + ev_lobpcg.data(), |
| 135 | + psi.data(), |
| 136 | + n, |
| 137 | + cg_tol, |
| 138 | + max_iter |
| 139 | + ); |
| 140 | + |
| 141 | + EXPECT_TRUE(converged) << "LOBPCG did not converge in " << max_iter << " iterations"; |
| 142 | + |
| 143 | + // Output LOBPCG eigenvalues for debug |
| 144 | + // std::cout << "LOBPCG computed eigenvalues (first 5): "; |
| 145 | + // for(int i=0; i<5 && i<nband; ++i) std::cout << ev_lobpcg[i] << " "; |
| 146 | + // std::cout << std::endl; |
| 147 | + |
| 148 | + // --------------------------------------------------------- |
| 149 | + // 3. Compare Results |
| 150 | + // --------------------------------------------------------- |
| 151 | + for(int i = 0; i < nband; ++i) { |
| 152 | + EXPECT_NEAR(ev_lobpcg[i], ev_lapack[i], check_tol) |
| 153 | + << "Mismatch at eigenvalue index " << i |
| 154 | + << " LAPACK: " << ev_lapack[i] << " LOBPCG: " << ev_lobpcg[i]; |
| 155 | + } |
| 156 | + // output anyway even if test passed, for debug |
| 157 | + std::cout << "Eigenvalues comparison (LOBPCG vs LAPACK):" << std::endl; |
| 158 | + std::cout << "LAPACK eigenvalues: "; |
| 159 | + for(int i=0; i<nband; ++i) { |
| 160 | + std::cout << "Index " << i << ": " << ev_lapack[i] << " vs " << ev_lapack[i] << std::endl; |
| 161 | + } |
| 162 | + std::cout << std::endl; |
| 163 | + std::cout << "LOBPCG eigenvalues: "; |
| 164 | + for(int i=0; i<nband; ++i) { |
| 165 | + std::cout << "Index " << i << ": " << ev_lobpcg[i] << " vs " << ev_lapack[i] << std::endl; |
| 166 | + } |
| 167 | + } |
| 168 | +}; |
| 169 | + |
| 170 | +TEST_F(DiagoLobpcgTest, CompareWithLapack) { |
| 171 | + int n = 100; |
| 172 | + int nband = 10; |
| 173 | + GenerateMatrix(n, 0); |
| 174 | + VerifyLobpcg(n, nband); |
| 175 | +} |
| 176 | + |
| 177 | +TEST_F(DiagoLobpcgTest, LargeScale) { |
| 178 | + int n = 200; |
| 179 | + int nband = 20; |
| 180 | + GenerateMatrix(n, 0); |
| 181 | + VerifyLobpcg(n, nband); |
| 182 | +} |
| 183 | + |
| 184 | +TEST_F(DiagoLobpcgTest, RandomMatrix) { |
| 185 | + int n = 50; |
| 186 | + int nband = 10; |
| 187 | + GenerateMatrix(n, 1); |
| 188 | + VerifyLobpcg(n, nband, 0.1, 1e-2); |
| 189 | +} |
0 commit comments