Merge pull request #1188 from deepmodeling/HSolver

dyzheng · web-flow · commit ebe893906c8c · 2022-08-25T19:46:24.000+08:00
Refactor: modified Operator::hPsi(), hpsi memory arranged outside
diff --git a/source/module_esolver/esolver_sdft_pw_tool.cpp b/source/module_esolver/esolver_sdft_pw_tool.cpp
@@ -246,22 +246,18 @@ void ESolver_SDFT_PW::sKG(const int nche_KG, const double fwhmin, const double w
         // this->phami->hPsi(j1psi.get_pointer(), j2psi.get_pointer(), ndim*totbands_per*npwx);
         // this->phami->hPsi(j1sfpsi.get_pointer(), j2sfpsi.get_pointer(), ndim*totbands_per*npwx);
         psi::Range allbands(1,0,0,totbands_per-1);
-        hamilt::Operator::hpsi_info info_psi0(&psi0, allbands);
-        const std::complex<double>* hpsi_out = std::get<0>(this->phami->ops->hPsi(info_psi0))->get_pointer();
-        ModuleBase::GlobalFunc::COPYARRAY(hpsi_out, hpsi0.get_pointer(), totbands_per*npwx);
+        hamilt::Operator<std::complex<double>>::hpsi_info info_psi0(&psi0, allbands, hpsi0.get_pointer());
+        this->phami->ops->hPsi(info_psi0);
 
-        hamilt::Operator::hpsi_info info_sfpsi0(&sfpsi0, allbands);
-        const std::complex<double>* hsfpsi_out = std::get<0>(this->phami->ops->hPsi(info_sfpsi0))->get_pointer();
-        ModuleBase::GlobalFunc::COPYARRAY(hsfpsi_out, hsfpsi0.get_pointer(), totbands_per*npwx);
+        hamilt::Operator<std::complex<double>>::hpsi_info info_sfpsi0(&sfpsi0, allbands, hsfpsi0.get_pointer());
+        this->phami->ops->hPsi(info_sfpsi0);
 
         psi::Range allndimbands(1,0,0,ndim*totbands_per-1);
-        hamilt::Operator::hpsi_info info_j1psi(&j1psi, allndimbands);
-        const std::complex<double>* hj1psi_out = std::get<0>(this->phami->ops->hPsi(info_j1psi))->get_pointer();
-        ModuleBase::GlobalFunc::COPYARRAY(hj1psi_out, j2psi.get_pointer(), ndim*totbands_per*npwx);
+        hamilt::Operator<std::complex<double>>::hpsi_info info_j1psi(&j1psi, allndimbands, j2psi.get_pointer());
+        this->phami->ops->hPsi(info_j1psi);
 
-        hamilt::Operator::hpsi_info info_j1sfpsi(&j1sfpsi, allndimbands);
-        const std::complex<double>* hj1sfpsi_out = std::get<0>(this->phami->ops->hPsi(info_j1sfpsi))->get_pointer();
-        ModuleBase::GlobalFunc::COPYARRAY(hj1sfpsi_out, j2sfpsi.get_pointer(), ndim*totbands_per*npwx);
+        hamilt::Operator<std::complex<double>>::hpsi_info info_j1sfpsi(&j1sfpsi, allndimbands, j2sfpsi.get_pointer());
+        this->phami->ops->hPsi(info_j1sfpsi);
 
         /*
         // stohchi.hchi_norm(psi0.get_pointer(), hpsi0.get_pointer(), totbands_per);
diff --git a/source/module_hamilt/hamilt.h b/source/module_hamilt/hamilt.h
@@ -34,7 +34,7 @@ class Hamilt
     int non_first_scf=0;
 
     // first node operator, add operations from each operators
-    Operator* ops = nullptr;
+    Operator<std::complex<double>>* ops = nullptr;
 };
 
 } // namespace hamilt
diff --git a/source/module_hamilt/hamilt_pw.cpp b/source/module_hamilt/hamilt_pw.cpp
@@ -25,7 +25,7 @@ HamiltPW::HamiltPW()
 
     if (GlobalV::T_IN_H)
     {
-        Operator* ekinetic = new Ekinetic<OperatorPW>( 
+        Operator<std::complex<double>>* ekinetic = new Ekinetic<OperatorPW>( 
             tpiba2, 
             gk2,
             GlobalC::wfcpw->npwk_max
@@ -41,7 +41,7 @@ HamiltPW::HamiltPW()
     }
     if (GlobalV::VL_IN_H)
     {
-        Operator* veff = new Veff<OperatorPW>(
+        Operator<std::complex<double>>* veff = new Veff<OperatorPW>(
             isk,
             &(GlobalC::pot.vr_eff),
             GlobalC::wfcpw
@@ -57,7 +57,7 @@ HamiltPW::HamiltPW()
     }
     if (GlobalV::VNL_IN_H)
     {
-        Operator* nonlocal = new Nonlocal<OperatorPW>(
+        Operator<std::complex<double>>* nonlocal = new Nonlocal<OperatorPW>(
             isk,
             &GlobalC::ppcell,
             &GlobalC::ucell
@@ -71,7 +71,7 @@ HamiltPW::HamiltPW()
             this->ops->add(nonlocal);
         }
     }
-    Operator* meta = new Meta<OperatorPW>(
+    Operator<std::complex<double>>* meta = new Meta<OperatorPW>(
         tpiba,
         isk,
         &GlobalC::pot.vofk,
diff --git a/source/module_hamilt/ks_pw/ekinetic_pw.cpp b/source/module_hamilt/ks_pw/ekinetic_pw.cpp
@@ -6,9 +6,6 @@
 namespace hamilt
 {
 
-template class Ekinetic<OperatorPW>;
-
-template<>
 Ekinetic<OperatorPW>::Ekinetic(
     double tpiba2_in,
     const double* gk2_in,
@@ -26,7 +23,6 @@ Ekinetic<OperatorPW>::Ekinetic(
     }
 }
 
-template<>
 void Ekinetic<OperatorPW>::act
 (
     const psi::Psi<std::complex<double>> *psi_in, 
diff --git a/source/module_hamilt/ks_pw/ekinetic_pw.h b/source/module_hamilt/ks_pw/ekinetic_pw.h
@@ -6,8 +6,16 @@
 namespace hamilt
 {
 
-template<class T>
-class Ekinetic : public T
+#ifndef __EKINETICTEMPLATE
+#define __EKINETICTEMPLATE
+
+template<class T> class Ekinetic : public T 
+{};
+
+#endif
+
+template<>
+class Ekinetic<OperatorPW> : public OperatorPW
 {
     public:
     Ekinetic(
diff --git a/source/module_hamilt/ks_pw/meta_pw.cpp b/source/module_hamilt/ks_pw/meta_pw.cpp
@@ -10,9 +10,6 @@
 namespace hamilt
 {
 
-template class Meta<OperatorPW>;
-
-template<>
 Meta<OperatorPW>::Meta(
     double tpiba_in, 
     const int* isk_in,
@@ -31,7 +28,6 @@ Meta<OperatorPW>::Meta(
     }
 }
 
-template<>
 void Meta<OperatorPW>::act
 (
     const psi::Psi<std::complex<double>> *psi_in, 
diff --git a/source/module_hamilt/ks_pw/meta_pw.h b/source/module_hamilt/ks_pw/meta_pw.h
@@ -8,8 +8,16 @@
 namespace hamilt
 {
 
-template<class T>
-class Meta : public T
+#ifndef __METATEMPLATE
+#define __METATEMPLATE
+
+template<class T> class Meta : public T
+{};
+
+#endif
+
+template<>
+class Meta<OperatorPW> : public OperatorPW
 {
     public:
     Meta(
diff --git a/source/module_hamilt/ks_pw/nonlocal_pw.cpp b/source/module_hamilt/ks_pw/nonlocal_pw.cpp
@@ -10,9 +10,6 @@
 namespace hamilt
 {
 
-template class Nonlocal<OperatorPW>;
-
-template<>
 Nonlocal<OperatorPW>::Nonlocal
 (
     const int* isk_in,
@@ -30,7 +27,6 @@ Nonlocal<OperatorPW>::Nonlocal
     }
 }
 
-template<>
 void Nonlocal<OperatorPW>::init(const int ik_in)
 {
     this->ik = ik_in;
@@ -49,7 +45,6 @@ void Nonlocal<OperatorPW>::init(const int ik_in)
 //--------------------------------------------------------------------------
 // this function sum up each non-local pseudopotential located on each atom,
 //--------------------------------------------------------------------------
-template<>
 void Nonlocal<OperatorPW>::add_nonlocal_pp(std::complex<double> *hpsi_in, const std::complex<double> *becp, const int m) const
 {
     ModuleBase::timer::tick("Nonlocal", "add_nonlocal_pp");
@@ -170,7 +165,6 @@ void Nonlocal<OperatorPW>::add_nonlocal_pp(std::complex<double> *hpsi_in, const
     return;
 }
 
-template<>
 void Nonlocal<OperatorPW>::act
 (
     const psi::Psi<std::complex<double>> *psi_in, 
diff --git a/source/module_hamilt/ks_pw/nonlocal_pw.h b/source/module_hamilt/ks_pw/nonlocal_pw.h
@@ -10,8 +10,16 @@
 namespace hamilt
 {
 
-template<class T>
-class Nonlocal : public T
+#ifndef __NONLOCALTEMPLATE
+#define __NONLOCALTEMPLATE
+
+template<class T> class Nonlocal : public T
+{};
+
+#endif
+
+template<>
+class Nonlocal<OperatorPW> : public OperatorPW
 {
     public:
     Nonlocal(
diff --git a/source/module_hamilt/ks_pw/operator_pw.h b/source/module_hamilt/ks_pw/operator_pw.h
@@ -6,17 +6,18 @@
 namespace hamilt
 {
 
-class OperatorPW : public Operator
+class OperatorPW : public Operator<std::complex<double>>
 {
     public:
     virtual ~OperatorPW(){};
     
     //in PW code, different operators donate hPsi independently
     //run this->act function for the first operator and run all act() for other nodes in chain table  
-    virtual hpsi_info hPsi(const hpsi_info& input)const
+    virtual hpsi_info hPsi(hpsi_info& input)const
     {
         ModuleBase::timer::tick("OperatorPW", "hPsi");
-        std::tuple<const std::complex<double>*, int> psi_info = std::get<0>(input)->to_range(std::get<1>(input));
+        auto psi_input = std::get<0>(input);
+        std::tuple<const std::complex<double>*, int> psi_info = psi_input->to_range(std::get<1>(input));
         int n_npwx = std::get<1>(psi_info); 
 
         std::complex<double> *tmhpsi = this->get_hpsi(input);
@@ -27,20 +28,25 @@ class OperatorPW : public Operator
             ModuleBase::WARNING_QUIT("OperatorPW", "please choose correct range of psi for hPsi()!");
         }
 
-        this->act(std::get<0>(input), n_npwx, tmpsi_in, tmhpsi);
+        this->act(psi_input, n_npwx, tmpsi_in, tmhpsi);
         OperatorPW* node((OperatorPW*)this->next_op);
         while(node != nullptr)
         {
-            node->act(std::get<0>(input), n_npwx, tmpsi_in, tmhpsi);
+            node->act(psi_input, n_npwx, tmpsi_in, tmhpsi);
             node = (OperatorPW*)(node->next_op);
         }
 
-        //during recursive call of hPsi, delete the input psi
-        if(this->recursive) delete std::get<0>(input);
-
         ModuleBase::timer::tick("OperatorPW", "hPsi");
         
-        return hpsi_info(this->hpsi, psi::Range(1, 0, 0, n_npwx/std::get<0>(input)->npol));
+        //if in_place, copy temporary hpsi to target hpsi_pointer, then delete hpsi and new a wrapper for return
+        std::complex<double>* hpsi_pointer = std::get<2>(input);
+        if(this->in_place)
+        {
+            ModuleBase::GlobalFunc::COPYARRAY(this->hpsi->get_pointer(), hpsi_pointer, this->hpsi->size());
+            delete this->hpsi;
+            this->hpsi = new psi::Psi<std::complex<double>>(hpsi_pointer, *psi_input, 1, n_npwx/psi_input->npol);
+        }      
+        return hpsi_info(this->hpsi, psi::Range(1, 0, 0, n_npwx/psi_input->npol), hpsi_pointer);
     }
     
     //main function which should be modified in Operator for PW base
diff --git a/source/module_hamilt/ks_pw/veff_pw.cpp b/source/module_hamilt/ks_pw/veff_pw.cpp
@@ -7,9 +7,6 @@
 namespace hamilt
 {
 
-template class Veff<OperatorPW>;
-
-template<>
 Veff<OperatorPW>::Veff(
     const int* isk_in,
     const ModuleBase::matrix* veff_in,
@@ -26,7 +23,6 @@ Veff<OperatorPW>::Veff(
     }
 }
 
-template<>
 void Veff<OperatorPW>::act
 (
     const psi::Psi<std::complex<double>> *psi_in, 
diff --git a/source/module_hamilt/ks_pw/veff_pw.h b/source/module_hamilt/ks_pw/veff_pw.h
@@ -8,8 +8,16 @@
 namespace hamilt
 {
 
-template<class T>
-class Veff : public T
+#ifndef __VEFFTEMPLATE
+#define __VEFFTEMPLATE
+
+template<class T> class Veff : public T 
+{};
+
+#endif
+
+template<>
+class Veff<OperatorPW> : public OperatorPW
 {
     public:
     Veff(
diff --git a/source/module_hamilt/operator.h b/source/module_hamilt/operator.h
@@ -4,10 +4,16 @@
 #include<complex>
 #include "module_psi/psi.h"
 #include "module_base/global_function.h"
+#include "module_base/tool_quit.h"
 
 namespace hamilt
 {
 
+// Basic class for operator module, 
+// it is designed for "O|psi>" and "<psi|O|psi>"
+// Operator "O" might have several different types, which should be calculated one by one.
+// In basic class , function add() is designed for combine all operators together with a chain. 
+template<typename T>
 class Operator
 {
     public:
@@ -25,11 +31,15 @@ class Operator
         } 
     }
 
-    typedef std::tuple<const psi::Psi<std::complex<double>>*, const psi::Range> hpsi_info;
-    virtual hpsi_info hPsi(const hpsi_info& input)const {return hpsi_info(nullptr, 0);}
-
-    virtual void act(std::complex<double> *hk_matrix)const {return;}
-    virtual void act(double *hk_matrix)const {return;}
+    //this is the core function for Operator
+    // do H|psi> from input |psi> , 
+    // output of hpsi would be first member of the returned tuple 
+    typedef std::tuple<const psi::Psi<T>*, const psi::Range, T*> hpsi_info;
+    virtual hpsi_info hPsi(hpsi_info& input)const 
+    {
+        ModuleBase::WARNING_QUIT("Operator::hPsi", "hPsi error!");
+        return hpsi_info(nullptr, 0, nullptr);
+    }
 
     virtual void init(const int ik_in)
     {
@@ -60,39 +70,48 @@ class Operator
     protected:
     int ik = 0;
 
-    mutable bool recursive = false;
+    mutable bool in_place = false;
 
     //calculation type, only different type can be in main chain table 
     int cal_type = 0;
     Operator* next_op = nullptr;
 
     //if this Operator is first node in chain table, hpsi would not be empty
-    mutable psi::Psi<std::complex<double>>* hpsi = nullptr;
+    mutable psi::Psi<T>* hpsi = nullptr;
 
-    std::complex<double>* get_hpsi(const hpsi_info& info)const
+    /*This function would analyze hpsi_info and choose how to arrange hpsi storage
+    In hpsi_info, if the third parameter hpsi_pointer is set, which indicates memory of hpsi is arranged by developer;
+    if hpsi_pointer is not set(nullptr), which indicates memory of hpsi is arranged by Operator, this case is rare. 
+    two cases would occurred:
+    1. hpsi_pointer != nullptr && psi_pointer == hpsi_pointer , psi would be replaced by hpsi, hpsi need a temporary memory
+    2. hpsi_pointer != nullptr && psi_pointer != hpsi_pointer , this is the commonly case 
+    */
+    T* get_hpsi(const hpsi_info& info)const
     {
         const int nbands_range = (std::get<1>(info).range_2 - std::get<1>(info).range_1 + 1);
-        //recursive call of hPsi, hpsi inputs as new psi, 
+        //in_place call of hPsi, hpsi inputs as new psi, 
         //create a new hpsi and delete old hpsi later
-        if(this->hpsi != std::get<0>(info) )
+        T* hpsi_pointer = std::get<2>(info);
+        const T* psi_pointer = std::get<0>(info)->get_pointer();
+        if(!hpsi_pointer)
         {
-            this->recursive = false;
-            if(this->hpsi != nullptr)
-            {
-                delete this->hpsi;
-            }
+            ModuleBase::WARNING_QUIT("Operator::hPsi", "hpsi_pointer can not be nullptr");
+        }
+        else if(hpsi_pointer == psi_pointer)
+        {
+            this->in_place = true;
+            this->hpsi = new psi::Psi<T>(std::get<0>(info)[0], 1, nbands_range);
         }
         else
         {
-            this->recursive = true;
+            this->in_place = false;
+            this->hpsi = new psi::Psi<T>(hpsi_pointer, std::get<0>(info)[0], 1, nbands_range);
         }
-        //create a new hpsi
-        this->hpsi = new psi::Psi<std::complex<double>>(std::get<0>(info)[0], 1, nbands_range);
         
-        std::complex<double>* pointer_hpsi = this->hpsi->get_pointer();
+        hpsi_pointer = this->hpsi->get_pointer();
         size_t total_hpsi_size = nbands_range * this->hpsi->get_nbasis();
-        ModuleBase::GlobalFunc::ZEROS(pointer_hpsi, total_hpsi_size);
-        return pointer_hpsi;
+        ModuleBase::GlobalFunc::ZEROS(hpsi_pointer, total_hpsi_size);
+        return hpsi_pointer;
     }
 };
 
diff --git a/source/module_hsolver/diago_cg.cpp b/source/module_hsolver/diago_cg.cpp
diff --git a/source/module_hsolver/diago_david.cpp b/source/module_hsolver/diago_david.cpp
diff --git a/source/module_hsolver/diago_iter_assist.cpp b/source/module_hsolver/diago_iter_assist.cpp
diff --git a/source/module_psi/psi.h b/source/module_psi/psi.h

Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,7 @@ HamiltPW::HamiltPW()`
`25`	`25`
`26`	`26`	`if (GlobalV::T_IN_H)`
`27`	`27`	`{`
`28`		`- Operator* ekinetic = new Ekinetic<OperatorPW>(`
	`28`	`+ Operator<std::complex<double>>* ekinetic = new Ekinetic<OperatorPW>(`
`29`	`29`	`tpiba2,`
`30`	`30`	`gk2,`
`31`	`31`	`GlobalC::wfcpw->npwk_max`
`@@ -41,7 +41,7 @@ HamiltPW::HamiltPW()`
`41`	`41`	`}`
`42`	`42`	`if (GlobalV::VL_IN_H)`
`43`	`43`	`{`
`44`		`- Operator* veff = new Veff<OperatorPW>(`
	`44`	`+ Operator<std::complex<double>>* veff = new Veff<OperatorPW>(`
`45`	`45`	`isk,`
`46`	`46`	`&(GlobalC::pot.vr_eff),`
`47`	`47`	`GlobalC::wfcpw`
`@@ -57,7 +57,7 @@ HamiltPW::HamiltPW()`
`57`	`57`	`}`
`58`	`58`	`if (GlobalV::VNL_IN_H)`
`59`	`59`	`{`
`60`		`- Operator* nonlocal = new Nonlocal<OperatorPW>(`
	`60`	`+ Operator<std::complex<double>>* nonlocal = new Nonlocal<OperatorPW>(`
`61`	`61`	`isk,`
`62`	`62`	`&GlobalC::ppcell,`
`63`	`63`	`&GlobalC::ucell`
`@@ -71,7 +71,7 @@ HamiltPW::HamiltPW()`
`71`	`71`	`this->ops->add(nonlocal);`
`72`	`72`	`}`
`73`	`73`	`}`
`74`		`- Operator* meta = new Meta<OperatorPW>(`
	`74`	`+ Operator<std::complex<double>>* meta = new Meta<OperatorPW>(`
`75`	`75`	`tpiba,`
`76`	`76`	`isk,`
`77`	`77`	`&GlobalC::pot.vofk,`
Original file line number	Diff line number	Diff line change
`@@ -6,9 +6,6 @@`
`6`	`6`	`namespace hamilt`
`7`	`7`	`{`
`8`	`8`
`9`		`-template class Ekinetic<OperatorPW>;`
`10`		`-`
`11`		`-template<>`
`12`	`9`	`Ekinetic<OperatorPW>::Ekinetic(`
`13`	`10`	`double tpiba2_in,`
`14`	`11`	`const double* gk2_in,`
`@@ -26,7 +23,6 @@ Ekinetic<OperatorPW>::Ekinetic(`
`26`	`23`	`}`
`27`	`24`	`}`
`28`	`25`
`29`		`-template<>`
`30`	`26`	`void Ekinetic<OperatorPW>::act`
`31`	`27`	`(`
`32`	`28`	`const psi::Psi<std::complex<double>> *psi_in,`
Original file line number	Diff line number	Diff line change
`@@ -10,9 +10,6 @@`
`10`	`10`	`namespace hamilt`
`11`	`11`	`{`
`12`	`12`
`13`		`-template class Meta<OperatorPW>;`
`14`		`-`
`15`		`-template<>`
`16`	`13`	`Meta<OperatorPW>::Meta(`
`17`	`14`	`double tpiba_in,`
`18`	`15`	`const int* isk_in,`
`@@ -31,7 +28,6 @@ Meta<OperatorPW>::Meta(`
`31`	`28`	`}`
`32`	`29`	`}`
`33`	`30`
`34`		`-template<>`
`35`	`31`	`void Meta<OperatorPW>::act`
`36`	`32`	`(`
`37`	`33`	`const psi::Psi<std::complex<double>> *psi_in,`