Add guard for hsolver rank-deficient cases (deepmodeling#7284)

* Add guard for hsolver rand-deficient cases

* Fix dimension of pw

Author: Cstandardlib

source/source_hsolver/diago_cg.cpp

@@ -542,6 +542,9 @@ void DiagoCG<T, Device>::schmit_orth(const int& m, const ct::Tensor& psi, const
                       << std::endl;
         }
         std::cout << " in DiagoCG, psi norm = " << psi_norm << std::endl;
+        std::cout << " This may be due to npwx < nbands: the number of plane waves is less than" << std::endl;
+        std::cout << " the number of bands, leading to a rank-deficient problem." << std::endl;
+        std::cout << " Please increase ecutwfc or reduce nbands." << std::endl;
         std::cout << " If you use GNU compiler, it may due to the zdotc is unavailable." << std::endl;
         ModuleBase::WARNING_QUIT("schmit_orth", "psi_norm <= 0.0");
     }

source/source_hsolver/diago_dav_subspace.cpp

@@ -4,6 +4,7 @@
 
 #include "source_base/module_device/device.h"
 #include "source_base/timer.h"
+#include "source_base/tool_quit.h"
 #include "source_base/kernels/math_kernel_op.h"
 #include "source_base/kernels/dsp/dsp_connector.h"
 // #include "source_base/module_container/ATen/kernels/lapack.h"
@@ -425,17 +426,54 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
                                 nbase,
                                 notconv,
                                 psi_norm);
+
+        // Check for zero norms (GPU path: copy norms from device to host)
+        Real* psi_norm_host = nullptr;
+        resmem_real_h_op()(psi_norm_host, notconv);
+        syncmem_var_d2h_op()(psi_norm_host, psi_norm, notconv);
+        for (int i = 0; i < notconv; i++)
+        {
+            if (psi_norm_host[i] <= 1.0e-12)
+            {
+                std::cout << "Diago_DavSubspace::cal_grad: psi_norm <= 0 for band " << i << std::endl;
+                std::cout << "This may be due to npwx < nbands: the number of plane waves is less than" << std::endl;
+                std::cout << "the number of bands, leading to a rank-deficient problem." << std::endl;
+                std::cout << "Please increase ecutwfc or reduce nbands." << std::endl;
+                delmem_real_h_op()(psi_norm_host);
+                delmem_real_op()(psi_norm);
+                ModuleBase::WARNING_QUIT("cal_grad", "psi_norm <= 0");
+            }
+        }
+        delmem_real_h_op()(psi_norm_host);
         delmem_real_op()(psi_norm);
     }
     else
 #endif
     {
         Real* psi_norm = nullptr;
+        resmem_real_h_op()(psi_norm, notconv);
+        setmem_real_h_op()(psi_norm, 0.0, notconv);
+
         normalize_op<T, Device>()(this->dim,
                                 psi_iter,
                                 nbase,
                                 notconv,
                                 psi_norm);
+
+        // Check for zero norms (CPU path)
+        for (int i = 0; i < notconv; i++)
+        {
+            if (psi_norm[i] <= 1.0e-12)
+            {
+                std::cout << "Diago_DavSubspace::cal_grad: psi_norm <= 0 for band " << i << std::endl;
+                std::cout << "This may be due to npwx < nbands: the number of plane waves is less than" << std::endl;
+                std::cout << "the number of bands, leading to a rank-deficient problem." << std::endl;
+                std::cout << "Please increase ecutwfc or reduce nbands." << std::endl;
+                delmem_real_h_op()(psi_norm);
+                ModuleBase::WARNING_QUIT("cal_grad", "psi_norm <= 0");
+            }
+        }
+        delmem_real_h_op()(psi_norm);
     }
 
     // update hpsi[:, nbase:nbase+notconv]

source/source_hsolver/diago_david.cpp

@@ -932,6 +932,9 @@ void DiagoDavid<T, Device>::SchmidtOrth(const int& dim,
     if (psi_norm < 1.0e-12)
     {
         std::cout << "DiagoDavid::SchmidtOrth:aborted for psi_norm <1.0e-12" << std::endl;
+        std::cout << "This may be due to npwx < nbands: the number of plane waves is less than" << std::endl;
+        std::cout << "the number of bands, leading to a rank-deficient problem." << std::endl;
+        std::cout << "Please increase ecutwfc or reduce nbands." << std::endl;
         std::cout << "nband = " << nband << std::endl;
         std::cout << "m = " << m << std::endl;
         exit(0);

source/source_hsolver/hsolver_pw.cpp

@@ -241,6 +241,29 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
 
     const int cur_nbasis = psi.get_current_nbas();
 
+    // Check for rank deficiency: the total number of plane waves (summed across
+    // all MPI processes) must be >= nbands. When npw_total < nbands, the basis is
+    // rank-deficient, leading to psi_norm <= 0 during Schmidt orthogonalization.
+    // Note: we sum cur_nbasis (local npw for this k-point) across the pool because
+    // psi.get_nbasis() gives the local storage dimension, not the total.
+    const int nbands = psi.get_nbands();
+    int npw_total = cur_nbasis;
+#ifdef __MPI
+    if (this->nproc_in_pool > 1)
+    {
+        MPI_Allreduce(&cur_nbasis, &npw_total, 1, MPI_INT, MPI_SUM, POOL_WORLD);
+    }
+#endif
+    if (npw_total < nbands)
+    {
+        const std::string msg = "npw_total < nbands (" + std::to_string(npw_total) + " < " + std::to_string(nbands)
+                            + "): the total number of plane waves across all MPI processes "
+                            + "is less than the number of bands, "
+                            + "which leads to a rank-deficient problem. "
+                            + "Please increase ecutwfc or reduce nbands.";
+        ModuleBase::WARNING_QUIT("HSolverPW::hamiltSolvePsiK", msg);
+    }
+
     // Shared matrix-blockvector operators used by all iterative solvers.
     auto hpsi_func = [hm, cur_nbasis](T* psi_in, T* hpsi_out, const int ld_psi, const int nvec) {
         auto psi_wrapper = psi::Psi<T, Device>(psi_in, 1, nvec, ld_psi, cur_nbasis);

source/source_hsolver/test/test_hsolver_pw.cpp

@@ -390,3 +390,19 @@ TEST_F(TestHSolverPW, SolveLcaoInPW) {
     EXPECT_DOUBLE_EQ(elecstate_test.ekb.c[0], 0.0);
     EXPECT_DOUBLE_EQ(elecstate_test.ekb.c[1], 0.0);
 }
+
+// Test that the program exits with an error when npwx < nbands,
+// which would cause rank deficiency and psi_norm <= 0 during diagonalization.
+TEST_F(TestHSolverPW, NpwxLessThanNbandsDeath)
+{
+    // Create psi with 5 bands but only 3 basis functions -> npwx=3 < nbands=5
+    psi_test_cd.resize(1, 5, 3);
+    std::vector<double> precond(3, 0.0);
+    std::vector<double> eigenvalues(5, 0.0);
+    // Expect death from WARNING_QUIT due to npwx < nbands
+    EXPECT_EXIT(
+        hs_d.hamiltSolvePsiK(&hamilt_test_d, psi_test_cd, precond, eigenvalues.data(), 1),
+        ::testing::ExitedWithCode(1),
+        ".*"
+    );
+}