Refactor: rename and privatize index mapping arrays in Parallel_2D; simplify deepks interfaces

source/module_basis/module_ao/ORB_control.cpp

@@ -213,7 +213,7 @@ void ORB_control::setup_2d_division(std::ofstream& ofs_running,
         ModuleBase::WARNING_QUIT("Parallel_Orbitals::set_global2local", "Check ks_solver.");
     }
     // (2) set the trace, then we can calculate the nnr.
-    // for 2d: calculate po.nloc first, then trace_loc_row and trace_loc_col
+    // for 2d: calculate po.nloc first, then global2local_row and global2local_col
     // for O(N): calculate the three together.
     this->ParaV.set_global2local(nlocal, nlocal, div_2d, ofs_running);
 }

source/module_basis/module_ao/parallel_2d.cpp

@@ -5,8 +5,8 @@ Parallel_2D::Parallel_2D()
 {}
 Parallel_2D::~Parallel_2D()
 {
-    delete[] trace_loc_row;
-    delete[] trace_loc_col;
+    delete[] global2local_row_;
+    delete[] global2local_col_;
 }
 
 void Parallel_2D::set_proc_dim(const int& dsize, bool mode /*= 0*/)
@@ -35,9 +35,9 @@ void Parallel_2D::set_proc_dim(const int& dsize, bool mode /*= 0*/)
 
 bool Parallel_2D::in_this_processor(const int& iw1_all, const int& iw2_all) const
 {
-    if (trace_loc_row[iw1_all] == -1)
+    if (global2local_row(iw1_all) == -1)
         return false;
-    else if (trace_loc_col[iw2_all] == -1)
+    else if (global2local_col(iw2_all) == -1)
         return false;
     return true;
 }
@@ -50,8 +50,8 @@ void Parallel_2D::set_global2local(const int& M_A, const int& N_A,
     this->init_global2local(M_A, N_A, ofs_running);
     if (!div_2d) // xiaohui add 2013-09-02
     {
-        for (int i = 0; i < M_A; i++) this->trace_loc_row[i] = i;
-        for (int i = 0; i < N_A; i++) this->trace_loc_col[i] = i;
+        for (int i = 0; i < M_A; i++) this->global2local_row_[i] = i;
+        for (int i = 0; i < N_A; i++) this->global2local_col_[i] = i;
         this->nrow = M_A;
         this->ncol = N_A;
         this->nloc = this->nrow * this->ncol;
@@ -61,19 +61,15 @@ void Parallel_2D::set_global2local(const int& M_A, const int& N_A,
         // ofs_running << " nrow=" << nrow << std::endl;
         for (int irow = 0; irow < this->nrow; irow++)
         {
-            int global_row = this->row_set[irow];
-            this->trace_loc_row[global_row] = irow;
-            // ofs_running << " global_row=" << global_row
-            // << " trace_loc_row=" << this->trace_loc_row[global_row] << std::endl;
+            int global_row = this->local2global_row(irow);
+            this->global2local_row_[global_row] = irow;
         }
 
         // ofs_running << " ncol=" << ncol << std::endl;
         for (int icol = 0; icol < this->ncol; icol++)
         {
-            int global_col = this->col_set[icol];
-            this->trace_loc_col[global_col] = icol;
-            // ofs_running << " global_col=" << global_col
-            // << " trace_loc_col=" << this->trace_loc_row[global_col] << std::endl;
+            int global_col = this->local2global_col(icol);
+            this->global2local_col_[global_col] = icol;
         }
     }
 
@@ -86,17 +82,17 @@ void Parallel_2D::init_global2local(const int& M_A, const int& N_A, std::ofstrea
     assert(M_A > 0);
     assert(N_A > 0);
 
-    delete[] this->trace_loc_row;
-    delete[] this->trace_loc_col;
+    delete[] this->global2local_row_;
+    delete[] this->global2local_col_;
 
-    ModuleBase::GlobalFunc::OUT(ofs_running, "trace_loc_row dimension", M_A);
-    ModuleBase::GlobalFunc::OUT(ofs_running, "trace_loc_col dimension", N_A);
+    ModuleBase::GlobalFunc::OUT(ofs_running, "global2local_row dimension", M_A);
+    ModuleBase::GlobalFunc::OUT(ofs_running, "global2local_col dimension", N_A);
 
-    this->trace_loc_row = new int[M_A];
-    this->trace_loc_col = new int[N_A];
+    this->global2local_row_ = new int[M_A];
+    this->global2local_col_ = new int[N_A];
 
-    for (int i = 0; i < M_A; i++) this->trace_loc_row[i] = -1;
-    for (int i = 0; i < N_A; i++) this->trace_loc_col[i] = -1;
+    for (int i = 0; i < M_A; i++) this->global2local_row_[i] = -1;
+    for (int i = 0; i < N_A; i++) this->global2local_col_[i] = -1;
 
     ModuleBase::Memory::record("trace_row_col", sizeof(int) * M_A);
     ModuleBase::Memory::record("trace_row_col", sizeof(int) * N_A);
@@ -242,17 +238,16 @@ int Parallel_2D::set_local2global(
 
     if (this->testpb)ModuleBase::GlobalFunc::OUT(ofs_running, "Local rows (including nb)", this->nrow);
 
-    // (5) row_set, it's a global index :
+    // (5) local2global_row, it's a global index :
     // save explicitly : every row in this processor
     // belongs to which row in the global matrix.
-    this->row_set.resize(this->nrow);
+    this->local2global_row_.resize(this->nrow);
     j = 0;
     for (int i = 0; i < row_b; i++)
     {
         for (int k = 0; k < nb && (coord[0] * nb + i * nb * dim[0] + k < M_A); k++, j++)
         {
-            this->row_set[j] = coord[0] * nb + i * nb * dim[0] + k;
-            // ofs_running << " j=" << j << " row_set=" << this->row_set[j] << std::endl;
+            this->local2global_row_[j] = coord[0] * nb + i * nb * dim[0] + k;
         }
     }
 
@@ -311,14 +306,14 @@ int Parallel_2D::set_local2global(
     //set nloc
     this->nloc = this->nrow * this->ncol;
 
-    this->col_set.resize(this->ncol);
+    this->local2global_col_.resize(this->ncol);
 
     j = 0;
     for (int i = 0; i < col_b; i++)
     {
         for (int k = 0; k < nb && (coord[1] * nb + i * nb * dim[1] + k < N_A); k++, j++)
         {
-            this->col_set[j] = coord[1] * nb + i * nb * dim[1] + k;
+            this->local2global_col_[j] = coord[1] * nb + i * nb * dim[1] + k;
         }
     }
     return 0;
@@ -330,9 +325,9 @@ void Parallel_2D::set_serial(const int& M_A, const int& N_A)
     this->nrow = M_A;
     this->ncol = N_A;
     this->nloc = this->nrow * this->ncol;
-    this->row_set.resize(this->nrow);
-    this->col_set.resize(this->ncol);
-    for (int i = 0; i < this->nrow; i++) this->row_set[i] = i;
-    for (int i = 0; i < this->ncol; i++) this->col_set[i] = i;
+    this->local2global_row_.resize(this->nrow);
+    this->local2global_col_.resize(this->ncol);
+    for (int i = 0; i < this->nrow; i++) this->local2global_row_[i] = i;
+    for (int i = 0; i < this->ncol; i++) this->local2global_col_[i] = i;
 }
 #endif

source/module_basis/module_ao/parallel_2d.h

@@ -15,14 +15,6 @@ class Parallel_2D
     Parallel_2D();
     ~Parallel_2D();
 
-    /// map from global index to local index
-    int* trace_loc_row = nullptr;
-    int* trace_loc_col = nullptr;
-
-    /// map from local index to global index
-    std::vector<int> row_set;				// Peize Lin change int* to vector 2022.08.03
-    std::vector<int> col_set;
-
     /// local size (nloc = nrow * ncol)
     int nrow;
     int ncol;
@@ -51,6 +43,30 @@ class Parallel_2D
     /// total number of elements of matrix in this processor
     int get_local_size()const { return this->nloc; };
 
+    /// get the local index of a global index (row)
+    int global2local_row(const int& igr) const
+    {
+        return this->global2local_row_[igr];
+    }
+
+    /// get the local index of a global index (col)
+    int global2local_col(const int& igc) const
+    {
+        return this->global2local_col_[igc];
+    }
+
+    /// get the global index of a local index (row)
+    int local2global_row(const int& ilr) const
+    {
+        return this->local2global_row_[ilr];
+    }
+
+    /// get the global index of a local index (col)
+    int local2global_col(const int& ilc) const
+    {
+        return this->local2global_col_[ilc];
+    }
+
     /// check whether an element is in this processor
     /// (check whether local-index > 0 )
     bool in_this_processor(const int& iw1_all, const int& iw2_all) const;
@@ -94,6 +110,15 @@ class Parallel_2D
         std::ofstream& ofs_running);
 
 protected:
+
+    /// map from global index to local index
+    int* global2local_row_ = nullptr;
+    int* global2local_col_ = nullptr;
+
+    /// map from local index to global index
+    std::vector<int> local2global_row_;				// Peize Lin change int* to vector 2022.08.03
+    std::vector<int> local2global_col_;
+
     /// set the map from local index to global index
     void init_global2local(const int& M_A/**< global row size*/,
         const int& N_A/**< global col size*/,

source/module_basis/module_ao/parallel_orbitals.cpp

@@ -41,7 +41,7 @@ void Parallel_Orbitals::set_atomic_trace(const int* iat2iwt, const int &nat, con
         //find the first row index of atom iat
         for(int i=0;i<max;i++)
         {
-            if(this->trace_loc_row[irow]!=-1)
+            if (this->global2local_row_[irow] != -1)
             {
                 this->atom_begin_row[iat] = irow;
                 break;
@@ -51,7 +51,7 @@ void Parallel_Orbitals::set_atomic_trace(const int* iat2iwt, const int &nat, con
         //find the first col index of atom iat
         for(int i=0;i<max;i++)
         {
-            if(this->trace_loc_col[icol]!=-1)
+            if (this->global2local_col_[icol] != -1)
             {
                 this->atom_begin_col[iat] = icol;
                 break;
@@ -136,7 +136,7 @@ int Parallel_Orbitals::set_nloc_wfc_Eij(
     std::ofstream& ofs_running,
     std::ofstream& ofs_warning)
 {
-    ModuleBase::TITLE("Parallel_Orbitals", "set_local2global");
+    ModuleBase::TITLE("Parallel_Orbitals", "set_nloc_wfc_Eij");
     // for wavefuncton , calculate nbands_loc
     int end_id;
     int block = N_A / nb;
@@ -154,7 +154,7 @@ int Parallel_Orbitals::set_nloc_wfc_Eij(
         }
         else
         {
-            ModuleBase::WARNING_QUIT("Parallel_Orbitals::set_local2global", "some processor has no bands-row-blocks.");
+            ModuleBase::WARNING_QUIT("Parallel_Orbitals::set_nloc_wfc_Eij", "some processor has no bands-row-blocks.");
         }
     }
     int col_b_bands = block / dim1;

source/module_basis/module_ao/test/parallel_2d_test.cpp

@@ -110,18 +110,21 @@ TEST_F(test_para2d, Divide2D)
 
                 //5. set_global2local
                 p2d.set_global2local(gr, gc, true, ofs_running);
-                auto sum_array = [](const int* arr, const int& size) -> int
+                auto sum_array = [&p2d](const int& gr, const int& gc) -> std::pair<int, int>
                     {
-                        int sum = 0;
-                        for (int i = 0; i < size; ++i)
-                            sum += arr[i];
-                        return sum;
+                        int sum_row = 0; int sum_col = 0;
+                        for (int i = 0; i < gr; ++i)
+                            sum_row += p2d.global2local_row(i);
+                        for (int i = 0; i < gc; ++i)
+                            sum_col += p2d.global2local_col(i);
+                        return { sum_row, sum_col };
                     };
-                EXPECT_EQ(sum_array(p2d.trace_loc_row, gr), lr * (lr - 1) / 2 - (gr - lr));
-                EXPECT_EQ(sum_array(p2d.trace_loc_col, gc), lc * (lc - 1) / 2 - (gc - lc));
+                std::pair<int, int> sumrc = sum_array(gr, gc);
+                EXPECT_EQ(std::get<0>(sumrc), lr * (lr - 1) / 2 - (gr - lr));
+                EXPECT_EQ(std::get<1>(sumrc), lc * (lc - 1) / 2 - (gc - lc));
                 for (int i = 0;i < lr;++i)
                     for (int j = 0;j < lc;++j)
-                        EXPECT_TRUE(p2d.in_this_processor(p2d.row_set[i], p2d.col_set[j]));
+                        EXPECT_TRUE(p2d.in_this_processor(p2d.local2global_row(i), p2d.local2global_col(j)));
             }
         }
     }
@@ -149,9 +152,9 @@ TEST_F(test_para2d, Serial)
         //3. set_global2local
         p2d.set_global2local(gr, gc, false, ofs_running);
         for (int i = 0;i < gr;++i)
-            EXPECT_EQ(p2d.trace_loc_row[i], i);
+            EXPECT_EQ(p2d.global2local_row(i), i);
         for (int i = 0;i < gc;++i)
-            EXPECT_EQ(p2d.trace_loc_col[i], i);
+            EXPECT_EQ(p2d.global2local_col(i), i);
     }
 }
 #endif

source/module_elecstate/cal_dm.h

@@ -33,13 +33,13 @@ inline void cal_dm(const Parallel_Orbitals* ParaV, const ModuleBase::matrix& wg,
         int ib_global = 0;
         for (int ib_local = 0; ib_local < nbands_local; ++ib_local)
         {
-            while (ib_local != ParaV->trace_loc_col[ib_global])
+            while (ib_local != ParaV->global2local_col(ib_global))
             {
                 ++ib_global;
                 if (ib_global >= wg.nc)
                 {
                     break;
-                    ModuleBase::WARNING_QUIT("ElecStateLCAO::cal_dm", "please check trace_loc_col!");
+                    ModuleBase::WARNING_QUIT("ElecStateLCAO::cal_dm", "please check global2local_col!");
                 }
             }
             if (ib_global >= wg.nc) continue;
@@ -91,13 +91,13 @@ inline void cal_dm(const Parallel_Orbitals* ParaV, const ModuleBase::matrix& wg,
         int ib_global = 0;
         for (int ib_local = 0; ib_local < nbands_local; ++ib_local)
         {
-            while (ib_local != ParaV->trace_loc_col[ib_global])
+            while (ib_local != ParaV->global2local_col(ib_global))
             {
                 ++ib_global;
                 if (ib_global >= wg.nc)
                 {
                     break;
-                    ModuleBase::WARNING_QUIT("ElecStateLCAO::cal_dm", "please check trace_loc_col!");
+                    ModuleBase::WARNING_QUIT("ElecStateLCAO::cal_dm", "please check global2local_col!");
                 }
             }
             if (ib_global >= wg.nc) continue;

source/module_esolver/esolver_ks_lcao.cpp

@@ -577,12 +577,12 @@ void ESolver_KS_LCAO::hamilt2density(int istep, int iter, double ethr)
         const Parallel_Orbitals* pv = this->LOWF.ParaV;
         if (GlobalV::GAMMA_ONLY_LOCAL)
         {
-            GlobalC::ld.cal_e_delta_band(this->LOC.dm_gamma, pv->trace_loc_row, pv->trace_loc_col, pv->nrow);
+            GlobalC::ld.cal_e_delta_band(this->LOC.dm_gamma);
         }
         else
         {
             GlobalC::ld
-                .cal_e_delta_band_k(this->LOC.dm_k, pv->trace_loc_row, pv->trace_loc_col, kv.nks, pv->nrow, pv->ncol);
+                .cal_e_delta_band_k(this->LOC.dm_k, kv.nks);
         }
     }
 #endif