fix UTs

deepmodeling · Nov 14, 2024 · fcff75e · fcff75e
1 parent 4f4f356
commit fcff75e
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 8 deletions.
diff --git a/source/module_hsolver/test/diago_david_float_test.cpp b/source/module_hsolver/test/diago_david_float_test.cpp
@@ -92,8 +92,9 @@ class DiagoDavPrepare
 
 		const int dim = phi.get_current_nbas() ;
 		const int nband = phi.get_nbands();
-		const int ld_psi =phi.get_nbasis();
-		hsolver::DiagoDavid<std::complex<float>> dav(precondition, nband, dim, order, false, comm_info);
+        const int ld_psi = phi.get_nbasis();
+        const hsolver::PreOP<std::complex<float>> pre_op(precondition, dim);
+        hsolver::DiagoDavid<std::complex<float>> dav(pre_op.get(), nband, dim, order, false, comm_info);
 
 		hsolver::DiagoIterAssist<std::complex<float>>::PW_DIAG_NMAX = maxiter;
 		hsolver::DiagoIterAssist<std::complex<float>>::PW_DIAG_THR = eps;

diff --git a/source/module_hsolver/test/diago_david_real_test.cpp b/source/module_hsolver/test/diago_david_real_test.cpp
@@ -92,7 +92,8 @@ class DiagoDavPrepare
 		const int dim = phi.get_current_nbas();
         const int nband = phi.get_nbands();
         const int ld_psi = phi.get_nbasis();
-        hsolver::DiagoDavid<double> dav(precondition, nband, dim, order, false, comm_info);
+        const hsolver::PreOP<double> pre_op(precondition, dim);
+        hsolver::DiagoDavid<double> dav(pre_op.get(), nband, dim, order, false, comm_info);
 
         hsolver::DiagoIterAssist<double>::PW_DIAG_NMAX = maxiter;
         hsolver::DiagoIterAssist<double>::PW_DIAG_THR = eps;

diff --git a/source/module_hsolver/test/diago_david_test.cpp b/source/module_hsolver/test/diago_david_test.cpp
@@ -91,8 +91,10 @@ class DiagoDavPrepare
 
 		const int dim = phi.get_current_nbas();
 		const int nband = phi.get_nbands();
-		const int ld_psi = phi.get_nbasis();
-		hsolver::DiagoDavid<std::complex<double>> dav(precondition, nband, dim, order, false, comm_info);
+        const int ld_psi = phi.get_nbasis();
+        const auto pre_func = [&precondition](std::complex<double>* ptr, const int& ld, const int& nvec)->void
+            {  hsolver::fvec::div_prevec(ptr, ld, nvec, precondition); };
+        hsolver::DiagoDavid<std::complex<double>> dav(pre_func, nband, dim, order, false, comm_info);
 
 		hsolver::DiagoIterAssist<std::complex<double>>::PW_DIAG_NMAX = maxiter;
 		hsolver::DiagoIterAssist<std::complex<double>>::PW_DIAG_THR = eps;

diff --git a/source/module_hsolver/test/hsolver_pw_sup.h b/source/module_hsolver/test/hsolver_pw_sup.h
@@ -1,4 +1,5 @@
 #include "module_basis/module_pw/pw_basis_k.h"
+#include "module_hsolver/precondition_funcs.h"
 
 namespace ModulePW {
 
@@ -121,15 +122,15 @@ template class DiagoCG<std::complex<float>, base_device::DEVICE_CPU>;
 template class DiagoCG<std::complex<double>, base_device::DEVICE_CPU>;
 
 template <typename T, typename Device>
-DiagoDavid<T, Device>::DiagoDavid(const Real* precondition_in,
+DiagoDavid<T, Device>::DiagoDavid(PreFunc&& precondition_in,
                                   const int nband_in,
                                   const int dim_in,
                                   const int david_ndim_in,
                                   const bool use_paw_in,
                                   const diag_comm_info& diag_comm_in)
-    : nband(nband_in), dim(dim_in), nbase_x(david_ndim_in * nband_in), david_ndim(david_ndim_in), use_paw(use_paw_in), diag_comm(diag_comm_in) {
+    : nband(nband_in), dim(dim_in), nbase_x(david_ndim_in* nband_in), david_ndim(david_ndim_in), use_paw(use_paw_in), diag_comm(diag_comm_in),
+    precondition(std::forward<PreFunc>(precondition_in)) {
     this->device = base_device::get_device_type<Device>(this->ctx);
-    this->precondition = precondition_in;
 
     test_david = 2;
     // 1: check which function is called and which step is executed