Change wrapper spsi_func in hsolver-dav (#5205)

Cstandardlib · pre-commit-ci-lite[bot] · web-flow · commit 5329628c0ab2 · 2024-10-10T22:45:13.000+08:00
* Change wrapper spsi_func in hsolver-dav

* Update spsi_func in hsolver_lrtd

* Update david test spsi_func

* [pre-commit.ci lite] apply automatic fixes

---------

Co-authored-by: pre-commit-ci-lite[bot] &lt;117423508+pre-commit-ci-lite[bot]@users.noreply.github.com&gt;
diff --git a/python/pyabacus/src/py_diago_david.hpp b/python/pyabacus/src/py_diago_david.hpp
@@ -132,7 +132,6 @@ class PyDiagoDavid
             const std::complex<double> *psi_in, 
             std::complex<double> *spsi_out, 
             const int nrow, 
-            const int npw, 
             const int nbands
         ) {
             syncmem_op()(this->ctx, this->ctx, spsi_out, psi_in, static_cast<size_t>(nbands * nrow));
diff --git a/source/module_hsolver/diago_david.cpp b/source/module_hsolver/diago_david.cpp
@@ -198,7 +198,7 @@ int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,
         else
         {
             // phm_in->sPsi(psi_in + m*ld_psi, &this->spsi[m * dim], dim, dim, 1);
-            spsi_func(psi_in + m*ld_psi,&this->spsi[m*dim],dim,dim,1);
+            spsi_func(psi_in + m*ld_psi,&this->spsi[m*dim],dim, 1);
         }
     }
     // begin SchmidtOrth
@@ -223,7 +223,7 @@ int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,
         else
         {
             // phm_in->sPsi(basis + dim*m, &this->spsi[m * dim], dim, dim, 1);
-            spsi_func(basis + dim*m, &this->spsi[m * dim], dim, dim, 1);
+            spsi_func(basis + dim*m, &this->spsi[m * dim], dim, 1);
         }
     }
 
@@ -554,7 +554,7 @@ void DiagoDavid<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
         else
         {
             // phm_in->sPsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, dim, 1);
-            spsi_func(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, dim, 1);
+            spsi_func(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, 1);
         }
     }
     // first nbase bands psi* dot notconv bands spsi to prepare lagrange_matrix
@@ -595,7 +595,7 @@ void DiagoDavid<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
         else
         {
             // phm_in->sPsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, dim, 1);
-            spsi_func(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, dim, 1);
+            spsi_func(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, 1);
         }
     }
     // calculate H|psi> for not convergence bands
diff --git a/source/module_hsolver/diago_david.h b/source/module_hsolver/diago_david.h
@@ -62,13 +62,10 @@ class DiagoDavid : public DiagH<T, Device>
      *
      * @param[in]   X       Pointer to the input blockvector.
      * @param[out] SX       Pointer to the output blockvector.
-     * @param[in] ld_spsi   Leading dimension of spsi. Dimension of SX: nbands * nrow.
-     * @param[in] ld_psi    Leading dimension of psi. Number of plane waves.
-     * @param[in] nbands    Number of vectors.
-     * 
-     * @note called like spsi(in, out, dim, dim, 1)
+     * @param[in] ld_psi    Leading dimension of psi and spsi. Dimension of X&SX: ld * nvec.
+     * @param[in] nvec      Number of vectors.
      */
-    using SPsiFunc = std::function<void(T*, T*, const int, const int, const int)>;
+    using SPsiFunc = std::function<void(T*, T*, const int, const int)>;
 
     int diag(
       const HPsiFunc& hpsi_func,  // function void hpsi(T*, T*, const int, const int) 
diff --git a/source/module_hsolver/hsolver_pw.cpp b/source/module_hsolver/hsolver_pw.cpp
@@ -511,17 +511,16 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
         };
 
         /// wrap spsi into lambda function, Matrix \times blockvector
-        /// spsi(X, SX, nrow, npw, nbands)
-        /// nrow is leading dimension of spsi, npw is leading dimension of psi, nbands is number of vecs
+        /// spsi(X, SX, ld, nvec)
+        /// ld is leading dimension of psi and spsi
         auto spsi_func = [hm](const T* psi_in, T* spsi_out,
-                               const int ld_spsi,  // Leading dimension of spsi. Dimension of SX: nbands * nrow.
-                               const int ld_psi,   // Leading dimension of psi. Number of plane waves.
+                               const int ld_psi,   // Leading dimension of psi and spsi.
                                const int nvec      // Number of vectors(bands)
                             ){
             ModuleBase::timer::tick("David", "spsi_func");
             // sPsi determines S=I or not by  PARAM.globalv.use_uspp inside
             // sPsi(psi, spsi, nrow, npw, nbands)
-            hm->sPsi(psi_in, spsi_out, ld_spsi, ld_psi, nvec);
+            hm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nvec);
             ModuleBase::timer::tick("David", "spsi_func");
         };
 
diff --git a/source/module_hsolver/test/diago_david_float_test.cpp b/source/module_hsolver/test/diago_david_float_test.cpp
@@ -46,8 +46,10 @@ void lapackEigen(int &npw, std::vector<std::complex<float>> &hm, float * e, bool
 	char tmp_c1 = 'V', tmp_c2 = 'U';
 	cheev_(&tmp_c1, &tmp_c2, &npw, tmp.data(), &npw, e, work2, &lwork, rwork, &info);
 	end = clock();
-	if(info) std::cout << "ERROR: Lapack solver, info=" << info <<std::endl;
-	if (outtime) std::cout<<"Lapack Run time: "<<(float)(end - start) / CLOCKS_PER_SEC<<" S"<<std::endl;
+	if(info) { std::cout << "ERROR: Lapack solver, info=" << info <<std::endl;
+}
+	if (outtime) { std::cout<<"Lapack Run time: "<<(float)(end - start) / CLOCKS_PER_SEC<<" S"<<std::endl;
+}
 
 	delete [] rwork;
 	delete [] work2;
@@ -74,7 +76,8 @@ class DiagoDavPrepare
 		//calculate eigenvalues by LAPACK;
 		float* e_lapack = new float[npw];
 		float* ev;
-		if(mypnum == 0) lapackEigen(npw, DIAGOTEST::hmatrix_f, e_lapack,DETAILINFO);
+		if(mypnum == 0) { lapackEigen(npw, DIAGOTEST::hmatrix_f, e_lapack,DETAILINFO);
+}
 
 		//do Diago_David::diag()
 		float* en = new float[npw];		
@@ -111,13 +114,13 @@ class DiagoDavPrepare
 					const int ld_psi, const int nvec)
                     {
                         auto psi_iter_wrapper = psi::Psi<std::complex<float>>(psi_in, 1, nvec, ld_psi, nullptr);
-                        psi::Range bands_range(1, 0, 0, nvec-1);
+                        psi::Range bands_range(true, 0, 0, nvec-1);
                         using hpsi_info = typename hamilt::Operator<std::complex<float>>::hpsi_info;
                         hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out);
                         phm->ops->hPsi(info);
                     };
-		auto spsi_func = [phm](const std::complex<float>* psi_in, std::complex<float>* spsi_out,const int nrow, const int npw,  const int nbands){
-			phm->sPsi(psi_in, spsi_out, nrow, npw, nbands);
+		auto spsi_func = [phm](const std::complex<float>* psi_in, std::complex<float>* spsi_out,const int ld_psi, const int nbands){
+			phm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nbands);
 		};
 		dav.diag(hpsi_func,spsi_func, ld_psi, phi.get_pointer(), en, eps, maxiter);
 
@@ -131,7 +134,8 @@ class DiagoDavPrepare
 
 		if(mypnum == 0)
 		{
-			if (DETAILINFO) std::cout<<"diag Run time: "<< use_time << std::endl;
+			if (DETAILINFO) { std::cout<<"diag Run time: "<< use_time << std::endl;
+}
 			for(int i=0;i<nband;i++)
 			{
 				EXPECT_NEAR(en[i],e_lapack[i],CONVTHRESHOLD);
@@ -148,8 +152,9 @@ class DiagoDavTest : public ::testing::TestWithParam<DiagoDavPrepare> {};
 TEST_P(DiagoDavTest,RandomHamilt)
 {
 	DiagoDavPrepare ddp = GetParam();
-	if (DETAILINFO&&ddp.mypnum==0) std::cout << "npw=" << ddp.npw << ", nband=" << ddp.nband << ", sparsity=" 
+	if (DETAILINFO&&ddp.mypnum==0) { std::cout << "npw=" << ddp.npw << ", nband=" << ddp.nband << ", sparsity=" 
 			  << ddp.sparsity << ", eps=" << ddp.eps << std::endl;
+}
 
     HPsi<std::complex<float>> hpsi(ddp.nband, ddp.npw, ddp.sparsity);
 	DIAGOTEST::hmatrix_f = hpsi.hamilt();
@@ -236,7 +241,8 @@ int main(int argc, char **argv)
 
     testing::InitGoogleTest(&argc, argv);
     ::testing::TestEventListeners &listeners = ::testing::UnitTest::GetInstance()->listeners();
-    if (myrank != 0) delete listeners.Release(listeners.default_result_printer());
+    if (myrank != 0) { delete listeners.Release(listeners.default_result_printer());
+}
 
     int result = RUN_ALL_TESTS();
     if (myrank == 0 && result != 0)
diff --git a/source/module_hsolver/test/diago_david_real_test.cpp b/source/module_hsolver/test/diago_david_real_test.cpp
@@ -118,8 +118,8 @@ class DiagoDavPrepare
                         hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out);
                         phm->ops->hPsi(info);
                     };
-        auto spsi_func = [phm](const double* psi_in, double* spsi_out,const int nrow, const int npw,  const int nbands){
-			phm->sPsi(psi_in, spsi_out, nrow, npw, nbands);
+        auto spsi_func = [phm](const double* psi_in, double* spsi_out,const int ld_psi, const int nbands){
+			phm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nbands);
 		};
         dav.diag(hpsi_func,spsi_func, ld_psi, phi.get_pointer(), en, eps, maxiter);
 
diff --git a/source/module_hsolver/test/diago_david_test.cpp b/source/module_hsolver/test/diago_david_test.cpp
@@ -118,8 +118,8 @@ class DiagoDavPrepare
                         hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out);
                         phm->ops->hPsi(info);
                     };
-		auto spsi_func = [phm](const std::complex<double>* psi_in, std::complex<double>* spsi_out,const int nrow, const int npw,  const int nbands){
-			phm->sPsi(psi_in, spsi_out, nrow, npw, nbands);
+		auto spsi_func = [phm](const std::complex<double>* psi_in, std::complex<double>* spsi_out,const int ld_psi, const int nbands){
+			phm->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nbands);
 		};
 		dav.diag(hpsi_func,spsi_func, ld_psi, phi.get_pointer(), en, eps, maxiter);
 
diff --git a/source/module_hsolver/test/hsolver_pw_sup.h b/source/module_hsolver/test/hsolver_pw_sup.h
@@ -154,7 +154,7 @@ DiagoDavid<T, Device>::~DiagoDavid() {
 
 template <typename T, typename Device>
 int DiagoDavid<T, Device>::diag(const std::function<void(T*, T*, const int, const int)>& hpsi_func,
-                                const std::function<void(T*, T*, const int, const int, const int)>& spsi_func,
+                                const std::function<void(T*, T*, const int, const int)>& spsi_func,
                                 const int ld_psi,
                                 T *psi_in,
                                 Real* eigenvalue_in,
diff --git a/source/module_lr/hsolver_lrtd.cpp b/source/module_lr/hsolver_lrtd.cpp
@@ -95,9 +95,10 @@ namespace LR
                         pHamilt->ops->hPsi(info);
                     };
                 auto spsi_func = [pHamilt](const T* psi_in, T* spsi_out,
-                               const int nrow, const int npw,  const int nbands){
+                               const int ld_psi, const int nbands)
+                {
                     // sPsi determines S=I or not by PARAM.globalv.use_uspp inside
-                    pHamilt->sPsi(psi_in, spsi_out, nrow, npw, nbands);
+                    pHamilt->sPsi(psi_in, spsi_out, ld_psi, ld_psi, nbands);
                 };
 
                 const int& dim = psi_k1_dav.get_nbasis();   //equals to leading dimension here

Original file line number	Diff line number	Diff line change
`@@ -198,7 +198,7 @@ int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,`
`198`	`198`	`else`
`199`	`199`	`{`
`200`	`200`	`// phm_in->sPsi(psi_in + mld_psi, &this->spsi[m dim], dim, dim, 1);`
`201`		`- spsi_func(psi_in + mld_psi,&this->spsi[mdim],dim,dim,1);`
	`201`	`+ spsi_func(psi_in + mld_psi,&this->spsi[mdim],dim, 1);`
`202`	`202`	`}`
`203`	`203`	`}`
`204`	`204`	`// begin SchmidtOrth`
`@@ -223,7 +223,7 @@ int DiagoDavid<T, Device>::diag_once(const HPsiFunc& hpsi_func,`
`223`	`223`	`else`
`224`	`224`	`{`
`225`	`225`	`// phm_in->sPsi(basis + dimm, &this->spsi[m dim], dim, dim, 1);`
`226`		`- spsi_func(basis + dimm, &this->spsi[m dim], dim, dim, 1);`
	`226`	`+ spsi_func(basis + dimm, &this->spsi[m dim], dim, 1);`
`227`	`227`	`}`
`228`	`228`	`}`
`229`	`229`
`@@ -554,7 +554,7 @@ void DiagoDavid<T, Device>::cal_grad(const HPsiFunc& hpsi_func,`
`554`	`554`	`else`
`555`	`555`	`{`
`556`	`556`	`// phm_in->sPsi(basis + dim(nbase + m), &spsi[(nbase + m) dim], dim, dim, 1);`
`557`		`- spsi_func(basis + dim(nbase + m), &spsi[(nbase + m) dim], dim, dim, 1);`
	`557`	`+ spsi_func(basis + dim(nbase + m), &spsi[(nbase + m) dim], dim, 1);`
`558`	`558`	`}`
`559`	`559`	`}`
`560`	`560`	`// first nbase bands psi* dot notconv bands spsi to prepare lagrange_matrix`
`@@ -595,7 +595,7 @@ void DiagoDavid<T, Device>::cal_grad(const HPsiFunc& hpsi_func,`
`595`	`595`	`else`
`596`	`596`	`{`
`597`	`597`	`// phm_in->sPsi(basis + dim(nbase + m), &spsi[(nbase + m) dim], dim, dim, 1);`
`598`		`- spsi_func(basis + dim(nbase + m), &spsi[(nbase + m) dim], dim, dim, 1);`
	`598`	`+ spsi_func(basis + dim(nbase + m), &spsi[(nbase + m) dim], dim, 1);`
`599`	`599`	`}`
`600`	`600`	`}`
`601`	`601`	`// calculate H\|psi> for not convergence bands`