refactor: use BaseCell in ESolver interface (#7700)

* refactor: use BaseCell in ESolver interface

* fix: align LCAO others override with BaseCell

* fix: link BaseCell in MD unit tests

---------

Co-authored-by: Fei Yang <2501213217@stu.pku.edu.cn>

Author: 19hello

source/source_cell/CMakeLists.txt

@@ -5,6 +5,7 @@ add_subdirectory(module_neighlist)
 add_library(
     cell
     OBJECT
+    base_cell.cpp
     atom_pseudo.cpp
     atom_spec.cpp
     pseudo.cpp

source/source_cell/base_cell.cpp

@@ -0,0 +1,12 @@
+#include "source_cell/base_cell.h"
+
+#include "source_base/tool_quit.h"
+
+void BaseCell::require_kind(const Kind& expected, const char* caller) const
+{
+    if (this->kind() != expected)
+    {
+        const char* required_cell = expected == Kind::unit_cell ? "UnitCell" : "MDCell";
+        ModuleBase::WARNING_QUIT(caller, std::string("This operation only supports ") + required_cell + ".");
+    }
+}

source/source_cell/base_cell.h

@@ -0,0 +1,58 @@
+#ifndef BASE_CELL_H
+#define BASE_CELL_H
+
+#include "source_base/matrix3.h"
+
+class BaseCell
+{
+public:
+    enum class Kind
+    {
+        unit_cell,
+        md_cell
+    };
+
+    virtual ~BaseCell() = default;
+
+    Kind kind() const
+    {
+        return get_kind();
+    }
+
+    int nat() const
+    {
+        return get_nat();
+    }
+
+    double lat0() const
+    {
+        return get_lat0();
+    }
+
+    double omega() const
+    {
+        return get_omega();
+    }
+
+    const ModuleBase::Matrix3& latvec() const
+    {
+        return get_latvec();
+    }
+
+    const ModuleBase::Matrix3& GT() const
+    {
+        return get_GT();
+    }
+
+    void require_kind(const Kind& expected, const char* caller) const;
+
+private:
+    virtual Kind get_kind() const = 0;
+    virtual int get_nat() const = 0;
+    virtual double get_lat0() const = 0;
+    virtual double get_omega() const = 0;
+    virtual const ModuleBase::Matrix3& get_latvec() const = 0;
+    virtual const ModuleBase::Matrix3& get_GT() const = 0;
+};
+
+#endif

source/source_cell/unitcell.h

@@ -7,12 +7,13 @@
 #include "source_cell/magnetism.h"
 #include "module_symmetry/symmetry.h"
 #include "source_cell/module_neighlist/atom_provider.h"
+#include "source_cell/base_cell.h"
 #include "source_cell/nonlocal_info_base.h"
 
 /**
  * @brief Provide the basic information about unitcell.
  */
-class UnitCell : public AtomProvider {
+class UnitCell : public AtomProvider, public BaseCell {
   public:
     double get_lat0() const override {
         return lat0;
@@ -287,6 +288,22 @@ class UnitCell : public AtomProvider {
     std::vector<ModuleBase::Vector3<double>> get_lambda() const;
     /// @brief get constrain for deltaspin
     std::vector<ModuleBase::Vector3<int>> get_constrain() const;
+
+  private:
+    Kind get_kind() const override
+    {
+        return Kind::unit_cell;
+    }
+
+    int get_nat() const override
+    {
+        return nat;
+    }
+
+    const ModuleBase::Matrix3& get_GT() const override
+    {
+        return GT;
+    }
 };
 
 #endif // unitcell class

source/source_esolver/esolver.cpp

@@ -318,5 +318,4 @@ ESolver* init_esolver(const Input_para& inp)
                                 + " line " + std::to_string(__LINE__));
 }
 
-
 } // namespace ModuleESolver

source/source_esolver/esolver.h

@@ -2,6 +2,7 @@
 #define ESOLVER_H
 
 #include "source_base/matrix.h"
+#include "source_cell/base_cell.h"
 #include "source_cell/unitcell.h"
 
 struct Input_para;
@@ -24,26 +25,26 @@ class ESolver
     }
 
     //! initialize the energy solver by using input parameters and cell modules
-    virtual void before_all_runners(UnitCell& ucell, const Input_para& inp) = 0;
+    virtual void before_all_runners(BaseCell& cell, const Input_para& inp) = 0;
 
     //! run energy solver
-    virtual void runner(UnitCell& cell, const int istep) = 0;
+    virtual void runner(BaseCell& cell, const int istep) = 0;
 
     //! perform post processing calculations
-    virtual void after_all_runners(UnitCell& ucell) = 0;
+    virtual void after_all_runners(BaseCell& cell) = 0;
 
     //! deal with exx and other calculation than scf/md/relax/cell-relax:
     //! such as nscf, get_wf and get_pchg
-    virtual void others(UnitCell& ucell, const int istep) {};
+    virtual void others(BaseCell&, const int) {}
 
     //! calculate total energy of a given system
     virtual double cal_energy() = 0;
 
     //! calcualte forces for the atoms in the given cell
-    virtual void cal_force(UnitCell& ucell, ModuleBase::matrix& force) = 0;
+    virtual void cal_force(BaseCell& cell, ModuleBase::matrix& force) = 0;
 
     //! calcualte stress of given cell
-    virtual void cal_stress(UnitCell& ucell, ModuleBase::matrix& stress) = 0;
+    virtual void cal_stress(BaseCell& cell, ModuleBase::matrix& stress) = 0;
 
     bool conv_esolver = true; // whether esolver is converged
 
@@ -54,7 +55,6 @@ class ESolver
  * @brief A subrutine called in init_esolver()
  *        This function returns type of ESolver
  *        Based on PARAM.inp.basis_type and PARAM.inp.esolver_type
- * 
  * @return [out] std::string The type of ESolver
  */
 std::string determine_type();
@@ -70,8 +70,6 @@ std::string determine_type();
  */
 ESolver* init_esolver(const Input_para& inp);
 
-
-
 } // namespace ModuleESolver
 
 #endif

source/source_esolver/esolver_dfpt_pw.cpp

@@ -7,75 +7,97 @@
 // ============================================================
 
 #include "esolver_dfpt_pw.h"
+
 #include "source_base/tool_quit.h"
 
-namespace ModuleESolver {
+namespace ModuleESolver
+{
 
-ESolver_DFPT_PW::ESolver_DFPT_PW() {
+ESolver_DFPT_PW::ESolver_DFPT_PW()
+{
     this->classname = "ESolver_DFPT_PW";
     this->basisname = "PW";
     gs_done_ = false;
     dfpt_ = nullptr;
 }
 
-ESolver_DFPT_PW::~ESolver_DFPT_PW() {
-    if (dfpt_ != nullptr) {
+ESolver_DFPT_PW::~ESolver_DFPT_PW()
+{
+    if (dfpt_ != nullptr)
+    {
         delete dfpt_;
         dfpt_ = nullptr;
     }
 }
 
-void ESolver_DFPT_PW::before_all_runners(UnitCell& ucell, const Input_para& inp) {
+void ESolver_DFPT_PW::before_all_runners(BaseCell& basecell, const Input_para& inp)
+{
+    basecell.require_kind(BaseCell::Kind::unit_cell, __FUNCTION__);
+    UnitCell& ucell = static_cast<UnitCell&>(basecell);
+
     ModuleBase::TITLE("ESolver_DFPT_PW", "before_all_runners");
-    
+
     ESolver_KS_PW<std::complex<double>, base_device::DEVICE_CPU>::before_all_runners(ucell, inp);
-    
+
     init_dfpt(ucell);
 }
 
-void ESolver_DFPT_PW::runner(UnitCell& ucell, const int istep) {
+void ESolver_DFPT_PW::runner(BaseCell& basecell, const int istep)
+{
+    basecell.require_kind(BaseCell::Kind::unit_cell, __FUNCTION__);
+    UnitCell& ucell = static_cast<UnitCell&>(basecell);
+
     ModuleBase::TITLE("ESolver_DFPT_PW", "runner");
-    
-    if (!gs_done_) {
+
+    if (!gs_done_)
+    {
         run_gs(ucell);
         gs_done_ = true;
     }
-    
-    if (dfpt_ != nullptr) {
+
+    if (dfpt_ != nullptr)
+    {
         dfpt_->run();
     }
-    
+
     run_post_process(ucell);
 }
 
-void ESolver_DFPT_PW::after_all_runners(UnitCell& ucell) {
+void ESolver_DFPT_PW::after_all_runners(BaseCell& basecell)
+{
+    basecell.require_kind(BaseCell::Kind::unit_cell, __FUNCTION__);
+    UnitCell& ucell = static_cast<UnitCell&>(basecell);
+
     ModuleBase::TITLE("ESolver_DFPT_PW", "after_all_runners");
-    
+
     ESolver_KS_PW<std::complex<double>, base_device::DEVICE_CPU>::after_all_runners(ucell);
 }
 
-void ESolver_DFPT_PW::run_gs(UnitCell& ucell) {
+void ESolver_DFPT_PW::run_gs(UnitCell& ucell)
+{
     ModuleBase::TITLE("ESolver_DFPT_PW", "run_gs");
-    
+
     ESolver_KS_PW<std::complex<double>, base_device::DEVICE_CPU>::runner(ucell, 0);
 }
 
-void ESolver_DFPT_PW::init_dfpt(UnitCell& ucell) {
+void ESolver_DFPT_PW::init_dfpt(UnitCell& ucell)
+{
     ModuleBase::TITLE("ESolver_DFPT_PW", "init_dfpt");
-    
+
     dfpt_ = new ModuleDFPT::DFPT_PW();
-    
-//    dfpt_->init(ucell, *this->stp.psi, this->pelec->nelec, PARAM.inp.ecutwfc);
-    
+
+    //    dfpt_->init(ucell, *this->stp.psi, this->pelec->nelec, PARAM.inp.ecutwfc);
+
     dfpt_->set_parameters("dfpt.in");
-    
+
     dfpt_->set_qmesh(1, 1, 1);
-    
+
     dfpt_->set_conv_thr(1e-8);
     dfpt_->set_max_iter(100);
 }
 
-void ESolver_DFPT_PW::run_post_process(UnitCell& ucell) {
+void ESolver_DFPT_PW::run_post_process(UnitCell& ucell)
+{
     ModuleBase::TITLE("ESolver_DFPT_PW", "run_post_process");
 }
 

source/source_esolver/esolver_dfpt_pw.h

@@ -12,26 +12,28 @@
 #include "esolver_ks_pw.h"
 #include "source_pw/module_dfpt/dfpt_pw.h"
 
-namespace ModuleESolver {
+namespace ModuleESolver
+{
 
-class ESolver_DFPT_PW : public ESolver_KS_PW<std::complex<double>, base_device::DEVICE_CPU> {
-public:
+class ESolver_DFPT_PW : public ESolver_KS_PW<std::complex<double>, base_device::DEVICE_CPU>
+{
+  public:
     ESolver_DFPT_PW();
     ~ESolver_DFPT_PW();
-    
-    void before_all_runners(UnitCell& ucell, const Input_para& inp) override;
-    void runner(UnitCell& ucell, const int istep) override;
-    void after_all_runners(UnitCell& ucell) override;
-    
-protected:
+
+    void before_all_runners(BaseCell& basecell, const Input_para& inp) override;
+    void runner(BaseCell& basecell, const int istep) override;
+    void after_all_runners(BaseCell& basecell) override;
+
+  protected:
     ModuleDFPT::DFPT_PW* dfpt_ = nullptr;
-    
+
     bool gs_done_ = false;
-    
+
     void run_gs(UnitCell& ucell);
-    
+
     void init_dfpt(UnitCell& ucell);
-    
+
     void run_post_process(UnitCell& ucell);
 };