From c4c99e1b8c1afed9de4c16a7541e1a47d42daa67 Mon Sep 17 00:00:00 2001
From: jinzx10 <jzx016@hotmail.com>
Date: Thu, 30 May 2024 09:05:04 +0800
Subject: [PATCH] Refactor: simplified Parallel_2d interface (#4237)

* clean up

* change to new interface

* fix missing header

* add missing header

* add missing header

* follow advice from review

* update BLACS test

* compromise to pyabacus compilation
---
 source/module_base/blacs_connector.h          |   6 +
 source/module_base/scalapack_connector.h      |  78 +----
 source/module_base/test/CMakeLists.txt        |   5 -
 .../module_base/test/blacs_connector_test.cpp |  53 ----
 .../module_base/test_parallel/CMakeLists.txt  |  14 +
 .../test_parallel/blacs_connector_test.cpp    | 109 +++++++
 .../test_parallel/blacs_connector_test.sh     |  15 +
 source/module_basis/module_ao/ORB_control.cpp |  27 +-
 source/module_basis/module_ao/parallel_2d.cpp | 275 +++++++-----------
 source/module_basis/module_ao/parallel_2d.h   | 175 ++++++-----
 .../module_ao/parallel_orbitals.cpp           |   6 +-
 .../module_ao/parallel_orbitals.h             |   3 +-
 .../module_ao/test/parallel_2d_test.cpp       |  75 ++---
 .../module_ao/test/parallel_orbitals_test.cpp |  32 +-
 .../module_dm/test/test_cal_dm_R.cpp          |  12 +-
 .../module_dm/test/test_cal_dmk_psi.cpp       |  12 +-
 .../module_dm/test/test_dm_R_init.cpp         |  12 +-
 .../module_dm/test/test_dm_constructor.cpp    |  12 +-
 .../module_dm/test/test_dm_io.cpp             |  12 +-
 source/module_esolver/io_npz.cpp              |  13 +-
 .../operator_lcao/test/test_T_NL_cd.cpp       |  12 +-
 .../operator_lcao/test/test_dftu.cpp          |  12 +-
 .../operator_lcao/test/test_ekineticnew.cpp   |  12 +-
 .../operator_lcao/test/test_nonlocalnew.cpp   |  12 +-
 .../operator_lcao/test/test_overlapnew.cpp    |  12 +-
 .../operator_lcao/test/test_overlapnew_cd.cpp |  12 +-
 .../test/test_sc_lambda_lcao.cpp              |   6 +-
 .../module_deepks/test/LCAO_deepks_test.h     |   2 +-
 .../test/LCAO_deepks_test_prep.cpp            |   3 +-
 .../test/cal_h_lambda_test.cpp                |   6 +-
 .../test/cal_mw_helper_test.cpp               |   6 +-
 .../test/spin_constrain_test.cpp              |   4 +-
 .../module_hcontainer/atom_pair.cpp           |   3 +-
 .../module_hcontainer/atom_pair.h             |   4 +-
 .../test/test_hcontainer.cpp                  |   4 +-
 .../test/test_hcontainer_complex.cpp          |   4 +-
 .../test/test_hcontainer_output.cpp           |   4 +-
 .../test/test_hcontainer_readCSR.cpp          |   2 +-
 .../module_hcontainer/test/test_transfer.cpp  |  14 +-
 .../module_tddft/bandenergy.h                 |   4 +-
 .../module_tddft/middle_hamilt.h              |   3 +-
 .../module_tddft/norm_psi.h                   |   3 +-
 .../module_tddft/propagator.h                 |   3 +-
 .../module_tddft/test/bandenergy_test.cpp     |   2 -
 .../module_tddft/test/middle_hamilt_test.cpp  |   2 -
 .../module_tddft/test/norm_psi_test.cpp       |   2 -
 .../module_tddft/test/propagator_test1.cpp    |   4 +-
 .../module_tddft/test/upsi_test1.cpp          |   2 -
 source/module_hamilt_lcao/module_tddft/upsi.h |   3 +-
 source/module_io/single_R_io.h                |   1 +
 source/module_io/test/single_R_io_test.cpp    |  11 +-
 .../test_serial/read_wfc_nao_test.cpp         |   2 -
 source/module_io/write_Vxc.hpp                |  12 +-
 53 files changed, 450 insertions(+), 684 deletions(-)
 delete mode 100644 source/module_base/test/blacs_connector_test.cpp
 create mode 100644 source/module_base/test_parallel/blacs_connector_test.cpp
 create mode 100644 source/module_base/test_parallel/blacs_connector_test.sh

diff --git a/source/module_base/blacs_connector.h b/source/module_base/blacs_connector.h
index cfb627a9e2..2d3baa4d5a 100644
--- a/source/module_base/blacs_connector.h
+++ b/source/module_base/blacs_connector.h
@@ -26,6 +26,9 @@
 //====================================8<----------------------------------------
 // blacs
     // Initialization
+#ifndef BLACS_CONNECTOR_H
+#define BLACS_CONNECTOR_H
+
 extern "C"
 {
 	void Cblacs_pinfo(int *myid, int *nprocs);
@@ -34,6 +37,7 @@ extern "C"
 		// Informational and Miscellaneous
 	void Cblacs_gridinfo(int icontxt, int* nprow, int *npcol, int *myprow, int *mypcol);
     void Cblacs_gridinit(int* icontxt, char* layout, int nprow, int npcol);
+    void Cblacs_gridexit(int* icontxt);
     int Cblacs_pnum(int icontxt, int prow, int pcol);
     void Cblacs_pcoord(int icontxt, int pnum, int *prow, int *pcol);
 	void Cblacs_exit(int icontxt);
@@ -45,4 +49,6 @@ extern "C"
 {
     int Csys2blacs_handle(MPI_Comm SysCtxt);
 }
+#endif // __MPI
+
 #endif
diff --git a/source/module_base/scalapack_connector.h b/source/module_base/scalapack_connector.h
index d8723d8880..6321b4e2c1 100644
--- a/source/module_base/scalapack_connector.h
+++ b/source/module_base/scalapack_connector.h
@@ -1,12 +1,12 @@
 #ifndef SCALAPACK_CONNECTOR_H
 #define SCALAPACK_CONNECTOR_H
 
+#ifdef __MPI
+
 #include <complex>
 
 extern "C"
 {
-	void blacs_gridinit_( int *ictxt, const char *order, const int *nprow, const int *npcol );
-	void blacs_gridinfo_( const int *ictxt, int *nprow, int *npcol, int *myprow, int *mypcol );
 	int numroc_( const int *n, const int *nb, const int *iproc, const int *srcproc, const int *nprocs );
 	void descinit_( 
 		int *desc, 
@@ -174,76 +174,6 @@ class ScalapackConnector
 	}
 };
 
-/*
-class ScalapackConnector
-{
-public:
-	static void transpose_desc( int desc_T[9], const int desc[9] )
-	{
-		desc_T[0] = desc[0];
-		desc_T[1] = desc[1];
-		desc_T[2] = desc[3];	desc_T[3] = desc[2];
-		desc_T[4] = desc[5];	desc_T[5] = desc[4];
-		desc_T[6] = desc[6];	desc_T[7] = desc[7];
-		desc_T[8] = desc[8];
-	}
-
-	static void blacs_gridinit( int &ictxt, const char order, const int nprow, const int npcol )
-	{
-		blacs_gridinit_(&ictxt, &order, &nprow, &npcol);
-	}
-	
-	static void blacs_gridinfo( const int &ictxt, int &nprow, int &npcol, int &myprow, int &mypcol )
-	{
-		blacs_gridinfo_( &ictxt, &nprow, &npcol, &myprow, &mypcol );
-	}
-	
-	static int numroc( const int n, const int nb, const int iproc, const int srcproc, const int nprocs )
-	{
-		return numroc_(&n, &nb, &iproc, &srcproc, &nprocs);
-	}
-	
-	static void descinit( 
-		int *desc, 
-		const int m, const int n, const int mb, const int nb, const int irsrc, const int icsrc, 
-		const int ictxt, const int lld, int &info )
-	{
-		descinit_(desc, &m, &n, &mb, &nb, &irsrc, &icsrc, &ictxt, &lld, &info);
-//		descinit_(desc, &n, &m, &nb, &mb, &irsrc, &icsrc, &ictxt, &lld, &info);
-	}
-	
-	// C = a * A.? * B.? + b * C
-	static void pgemm(
-		const char transa, const char transb,
-		const int M, const int N, const int K,
-		const double alpha,
-		const double *A, const int IA, const int JA, const int *DESCA,
-		const double *B, const int IB, const int JB, const int *DESCB,
-		const double beta,
-		double *C, const int IC, const int JC, const int *DESCC)
-	{
-//		int DESCA_T[9], DESCB_T[9], DESCC_T[9];
-//		transpose_desc( DESCA_T, DESCA );
-//		transpose_desc( DESCB_T, DESCB );
-//		transpose_desc( DESCC_T, DESCC );
-//		pdgemm_(
-//			&transb, &transa,
-//			&N, &M, &K,
-//			&alpha,
-//			B, &JB, &IB, DESCB_T,
-//			A, &JA, &IA, DESCA_T,
-//			&beta,
-//			C, &JC, &IC, DESCC_T);
-		pdgemm_(
-			&transa, &transb,
-			&M, &N, &K,
-			&alpha,
-			A, &JA, &IA, DESCA,
-			B, &JB, &IB, DESCB,
-			&beta,
-			C, &JC, &IC, DESCC);
-	}
-};
-*/
+#endif // __MPI
 
-#endif
\ No newline at end of file
+#endif
diff --git a/source/module_base/test/CMakeLists.txt b/source/module_base/test/CMakeLists.txt
index 68556121d4..c4fc2c9a07 100644
--- a/source/module_base/test/CMakeLists.txt
+++ b/source/module_base/test/CMakeLists.txt
@@ -9,11 +9,6 @@ AddTest(
   TARGET base_atom_in
   SOURCES atom_in_test.cpp
 )
-AddTest(
-  TARGET base_blacs_connector
-  LIBS ${math_libs}
-  SOURCES blacs_connector_test.cpp
-)
 AddTest(
   TARGET base_timer
   SOURCES timer_test.cpp ../timer.cpp  ../global_variable.cpp
diff --git a/source/module_base/test/blacs_connector_test.cpp b/source/module_base/test/blacs_connector_test.cpp
deleted file mode 100644
index e1313d4e41..0000000000
--- a/source/module_base/test/blacs_connector_test.cpp
+++ /dev/null
@@ -1,53 +0,0 @@
-#include "../blacs_connector.h"
-
-#include <mpi.h>
-
-#include "gtest/gtest.h"
-
-/************************************************
- *  unit test of functions in blacs_connector.h
- ***********************************************/
-
-/**
- * - Tested Function
- *   - Cblacs_gridinit
- *     - Initializes a grid of processors with a given number of rows and columns.
- *        The function creates a cartesian topology of all the processors initialized
- *        by the BLS library. In this topology, each processor is identified by its
- *        coordinates (row, col) in the grid.
- */
-
-TEST(blacs_connector, Cblacs_gridinit)
-{
-    int icontxt = 0;
-    char layout[] = "ROW";
-    int nprow = 1;
-    int npcol = 1;
-
-    int myid = 0, nprocs = 0;
-    Cblacs_pinfo(&myid, &nprocs);
-    Cblacs_get(-1, 0, &icontxt);
-
-    // Call the Cblacs_gridinit() function
-    Cblacs_gridinit(&icontxt, layout, nprow, npcol);
-
-    // Check if the grid context handle is created successfully
-    EXPECT_EQ(icontxt, 0);
-}
-
-int main(int argc, char** argv)
-{
-    int myrank = 0;
-    int mysize = 0;
-
-    MPI_Init(&argc, &argv);
-    MPI_Comm_size(MPI_COMM_WORLD, &mysize);
-    MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
-
-    testing::InitGoogleTest(&argc, argv);
-
-    int result = 0;
-    result = RUN_ALL_TESTS();
-    MPI_Finalize();
-    return 0;
-}
diff --git a/source/module_base/test_parallel/CMakeLists.txt b/source/module_base/test_parallel/CMakeLists.txt
index 3ef9779fa6..9e05105183 100644
--- a/source/module_base/test_parallel/CMakeLists.txt
+++ b/source/module_base/test_parallel/CMakeLists.txt
@@ -33,3 +33,17 @@ add_test(NAME base_parallel_reduce_test
       COMMAND ${BASH} parallel_reduce_test.sh
       WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}
 )
+
+if(ENABLE_LCAO)
+AddTest(
+  TARGET blacs_connector
+  LIBS MPI::MPI_CXX ScaLAPACK::ScaLAPACK
+  SOURCES blacs_connector_test.cpp
+)
+install(FILES blacs_connector_test.sh DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+add_test(NAME blacs_connector_test
+      COMMAND ${BASH} blacs_connector_test.sh
+      WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}
+)
+endif()
+
diff --git a/source/module_base/test_parallel/blacs_connector_test.cpp b/source/module_base/test_parallel/blacs_connector_test.cpp
new file mode 100644
index 0000000000..59fd348fc9
--- /dev/null
+++ b/source/module_base/test_parallel/blacs_connector_test.cpp
@@ -0,0 +1,109 @@
+#ifdef __MPI
+
+#include "../blacs_connector.h"
+#include <mpi.h>
+#include "gtest/gtest.h"
+
+/************************************************
+ *  unit test of functions in blacs_connector.h
+ ***********************************************/
+
+/**
+ * - Tested Function
+ *   - Cblacs_gridinit
+ *     - Initializes a grid of processors with a given number of rows and columns.
+ *        The function creates a cartesian topology of all the processors initialized
+ *        by the BLS library. In this topology, each processor is identified by its
+ *        coordinates (row, col) in the grid.
+ */
+
+class BLACSTest: public testing::Test
+{
+protected:
+    void SetUp();
+
+    int rank = 0;
+    int nprocs = 0;
+    char layout = 'R';
+
+    // number of rows and columns in the process grid
+    int nprow = 0;
+    int npcol = 0;
+
+    // process coordinate
+    int iprow = -1;
+    int ipcol = -1;
+};
+
+void BLACSTest::SetUp()
+{
+    Cblacs_pinfo(&rank, &nprocs);
+}
+
+
+TEST_F(BLACSTest, WorldGrid)
+{
+    // generate a grid of size 1 x nproc
+    nprow = 1;
+    npcol = nprocs;
+
+    int ictxt_row = Csys2blacs_handle(MPI_COMM_WORLD);
+    Cblacs_gridinit(&ictxt_row, &layout, nprow, npcol);
+    Cblacs_gridinfo(ictxt_row, &nprow, &npcol, &iprow, &ipcol);
+
+    EXPECT_EQ(iprow, 0);
+    EXPECT_EQ(ipcol, rank);
+
+    // generate a grid of size nproc x 1
+    nprow = nprocs;
+    npcol = 1;
+
+    int ictxt_col = Csys2blacs_handle(MPI_COMM_WORLD);
+    Cblacs_gridinit(&ictxt_col, &layout, nprow, npcol);
+    Cblacs_gridinfo(ictxt_col, &nprow, &npcol, &iprow, &ipcol);
+
+    EXPECT_EQ(iprow, rank);
+    EXPECT_EQ(ipcol, 0);
+
+
+    // two BLACS grids should have difference context index
+    EXPECT_NE(ictxt_row, ictxt_col);
+}
+
+TEST_F(BLACSTest, SplitGrid)
+{
+    // this test create BLACS grids based on a disjoint communicator
+
+    const int n_blacs = 2;
+    int rank_sub = -1;
+    int nprocs_sub = 0;
+
+    // sub communicators are divided based on odd / even ranks
+    MPI_Comm comm_sub;
+    MPI_Comm_split(MPI_COMM_WORLD, rank % n_blacs, rank, &comm_sub);
+    MPI_Comm_rank(comm_sub, &rank_sub);
+    MPI_Comm_size(comm_sub, &nprocs_sub);
+
+    int ctxt_sub = Csys2blacs_handle(comm_sub);
+
+    nprow = 1, npcol = nprocs_sub; // row-like grids
+    Cblacs_gridinit(&ctxt_sub, &layout, nprow, npcol);
+    Cblacs_gridinfo(ctxt_sub, &nprow, &npcol, &iprow, &ipcol);
+
+    // verifies that the BLACS grid is created based on comm_sub instead of MPI_COMM_WORLD
+    EXPECT_EQ(iprow, 0);
+    EXPECT_EQ(ipcol, rank_sub);
+}
+
+int main(int argc, char** argv)
+{
+    MPI_Init(&argc, &argv);
+    testing::InitGoogleTest(&argc, argv);
+
+    int result = RUN_ALL_TESTS();
+
+    MPI_Finalize();
+
+    return result;
+}
+#endif
diff --git a/source/module_base/test_parallel/blacs_connector_test.sh b/source/module_base/test_parallel/blacs_connector_test.sh
new file mode 100644
index 0000000000..eca32b5bed
--- /dev/null
+++ b/source/module_base/test_parallel/blacs_connector_test.sh
@@ -0,0 +1,15 @@
+#!/bin/bash -e
+
+np=`cat /proc/cpuinfo | grep "cpu cores" | uniq| awk '{print $NF}'`
+echo "nprocs in this machine is $np"
+
+for i in {8..2};
+do
+    if [[ $i -gt $np ]];then
+        continue
+    fi
+    echo "TEST in parallel, nprocs=$i"
+    mpirun -np $i ./blacs_connector
+    break    
+done
+
diff --git a/source/module_basis/module_ao/ORB_control.cpp b/source/module_basis/module_ao/ORB_control.cpp
index c3a0bb41d2..82defaecb6 100644
--- a/source/module_basis/module_ao/ORB_control.cpp
+++ b/source/module_basis/module_ao/ORB_control.cpp
@@ -210,12 +210,8 @@ void ORB_control::setup_2d_division(std::ofstream& ofs_running,
     else
     {
         std::cout << " Parallel Orbial, DIAGO_TYPE = " << ks_solver << std::endl;
-        ModuleBase::WARNING_QUIT("Parallel_Orbitals::set_global2local", "Check ks_solver.");
+        ModuleBase::WARNING_QUIT("Parallel_Orbitals::setup_2d_division", "Check ks_solver.");
     }
-    // (2) set the trace, then we can calculate the nnr.
-    // 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);
 }
 
 
@@ -336,12 +332,6 @@ void ORB_control::divide_HS_2d(
     if (dcolor != 0)
         return; // mohan add 2012-01-13
 
-    // get the 2D index of computer.
-    pv->dim0 = (int)sqrt((double)dsize); // mohan update 2012/01/13
-    // while (GlobalV::NPROC_IN_POOL%dim0!=0)
-
-    pv->set_proc_dim(dsize);
-
     if (pv->testpb)
         ModuleBase::GlobalFunc::OUT(ofs_running, "dim0", pv->dim0);
     if (pv->testpb)
@@ -352,32 +342,27 @@ void ORB_control::divide_HS_2d(
 #ifdef __DEBUG
 assert(nb2d > 0);
 #endif
-    pv->set_block_size(nb2d); // mohan add 2010-06-28
 
     ModuleBase::GlobalFunc::OUT(ofs_running, "nb2d", pv->get_block_size());
 
     this->set_parameters(ofs_running, ofs_warning);
 
-    // call mpi_creat_cart
-    pv->mpi_create_cart(DIAG_WORLD);
-
-    int try_nb = pv->set_local2global(nlocal, nlocal, ofs_running, ofs_warning);
-    try_nb = pv->set_nloc_wfc_Eij(nbands, ofs_running, ofs_warning);
-    if (try_nb == 1)
+    int try_nb = pv->init(nlocal, nlocal, nb2d, DIAG_WORLD);
+    try_nb += pv->set_nloc_wfc_Eij(nbands, ofs_running, ofs_warning);
+    if (try_nb != 0)
     {
         ofs_running << " parameter nb2d is too large: nb2d = " << pv->get_block_size() << std::endl;
         ofs_running << " reset nb2d to value 1, this set would make the program keep working but maybe get slower "
                        "during diagonalization."
                     << std::endl;
-        pv->set_block_size(1);
-        try_nb = pv->set_local2global(nlocal, nlocal, ofs_running, ofs_warning);
+
+        pv->set(nlocal, nlocal, 1, pv->comm_2D, pv->blacs_ctxt);
         try_nb = pv->set_nloc_wfc_Eij(nbands, ofs_running, ofs_warning);
     }
 
     // init blacs context for genelpa
     if (ks_solver == "genelpa" || ks_solver == "scalapack_gvx" || ks_solver == "cusolver" || ks_solver == "cg_in_lcao" || ks_solver == "pexsi")
     {
-        pv->set_desc(nlocal, nlocal, pv->nrow);
         pv->set_desc_wfc_Eij(nlocal, nbands, pv->nrow);
     }
 #else // single processor used.
diff --git a/source/module_basis/module_ao/parallel_2d.cpp b/source/module_basis/module_ao/parallel_2d.cpp
index a75d56b2d5..4ac2bfefe5 100644
--- a/source/module_basis/module_ao/parallel_2d.cpp
+++ b/source/module_basis/module_ao/parallel_2d.cpp
@@ -1,209 +1,130 @@
 #include "parallel_2d.h"
-#include "module_base/memory.h"
 
-Parallel_2D::Parallel_2D()
-{}
-Parallel_2D::~Parallel_2D()
-{
-    delete[] global2local_row_;
-    delete[] global2local_col_;
-}
-
-void Parallel_2D::set_proc_dim(const int& dsize, bool mode /*= 0*/)
-{
-    // default mode = 0: dim0 <= dim1
-    this->dim0 = (int)sqrt((double)dsize);
-    while (dsize % this->dim0 != 0)
-    {
-        this->dim0 = this->dim0 - 1;
-    }
-    assert(this->dim0 > 0);
-    this->dim1 = dsize / this->dim0;
+#include <numeric>
+#include <cassert>
 
-    if (mode) { // mode = 1: dim0 >= dim1
-        std::swap(this->dim0, this->dim1);
-    }
-}
+#include "module_base/blacs_connector.h"
+#include "module_base/scalapack_connector.h"
 
-bool Parallel_2D::in_this_processor(const int& iw1_all, const int& iw2_all) const
+bool Parallel_2D::in_this_processor(const int iw1_all, const int iw2_all) const
 {
     return global2local_row(iw1_all) != -1 && global2local_col(iw2_all) != -1;
 }
 
-void Parallel_2D::set_global2local(const int& M_A, const int& N_A,
-    const bool& div_2d, std::ofstream& ofs_running)
-{
-    ModuleBase::TITLE("Parallel_2D", "set_global2local");
-
-    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->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;
-    }
-    else
-    {
-        // ofs_running << " nrow=" << nrow << std::endl;
-        for (int irow = 0; irow < this->nrow; irow++)
-        {
-            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->local2global_col(icol);
-            this->global2local_col_[global_col] = icol;
-        }
-    }
-
-    return;
-}
-
-void Parallel_2D::init_global2local(const int& M_A, const int& N_A, std::ofstream& ofs_running)
-{
-    ModuleBase::TITLE("Parallel_2D", "init_global2local");
-    assert(M_A > 0);
-    assert(N_A > 0);
-
-    delete[] this->global2local_row_;
-    delete[] this->global2local_col_;
-
-    ModuleBase::GlobalFunc::OUT(ofs_running, "global2local_row dimension", M_A);
-    ModuleBase::GlobalFunc::OUT(ofs_running, "global2local_col dimension", 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->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);
-    return;
-}
-
 #ifdef __MPI
-extern "C"
+void Parallel_2D::_init_proc_grid(const MPI_Comm comm, const bool mode)
 {
-#include "module_base/blacs_connector.h"
-#include "module_base/scalapack_connector.h"
+    // determine the number of rows and columns of the process grid
+    // by factorizing n = p * q such that p, q are closest and p <= q
+    int num_proc = 0;
+    MPI_Comm_size(comm, &num_proc);
+    dim0 = static_cast<int>(std::sqrt(num_proc + 0.5));
+    while (dim1 = num_proc / dim0, dim0 * dim1 != num_proc) { --dim0; }
+
+    if (mode) { std::swap(dim0, dim1); }
+
+    // create a 2D Cartesian MPI communicator (row-major by default)
+    int period[2] = {1, 1};
+    int dim[2] = {dim0, dim1};
+    const int reorder = 0;
+    MPI_Cart_create(comm, 2, dim, period, reorder, &comm_2D);
+    MPI_Cart_get(comm_2D, 2, dim, period, coord);
+
+    // initialize the BLACS grid accordingly
+    blacs_ctxt = Csys2blacs_handle(comm_2D);
+    char order = 'R'; // row-major
+    Cblacs_gridinit(&blacs_ctxt, &order, dim0, dim1);
+
+    // TODO Currently MPI and BLACS are made to have the same Cartesian grid.
+    // In theory, however, BLACS would split any given communicator to create
+    // new ones for its own purpose when initializing the process grid, so it
+    // might be unnecessary to create an MPI communicator with Cartesian topology.
+    // ***This needs to be verified***
 }
 
-// FIXME In theory BLACS would split the given communicator to get some new
-// ones for its own purpose when initializing the process grid, so there might
-// be unnecessary to create a Cartesian MPI communicator in advance.
-// ***This needs to be verified***
 
-void Parallel_2D::mpi_create_cart(const MPI_Comm& diag_world)
+void Parallel_2D::_set_dist_info(const int mg, const int ng, const int nb)
 {
-    ModuleBase::TITLE("Parallel_2D", "mpi_create_cart");
-#ifdef __DEBUG
-    assert(this->comm_2D != MPI_COMM_NULL);
-    assert(this->dim0 > 0 && this->dim1 > 0);
-#endif
-    // the matrix is divided as ( dim0 * dim1 )
-    int period[2] = { 1,1 };
-    int dim[2] = { this->dim0, this->dim1 };
-    int reorder = 0;
-    MPI_Cart_create(diag_world, 2, dim, period, reorder, &this->comm_2D);
-    return;
-}
+    this->nb = nb;
 
-void Parallel_2D::set_desc(const int& gr, const int& gc, const int& lld, bool first_time)
-{
-    ModuleBase::TITLE("Parallel_2D", "set_desc");
-#ifdef __DEBUG
-    assert(this->comm_2D != MPI_COMM_NULL);
-    assert(gr > 0 && gc > 0 && lld > 0);
-    assert(this->nb > 0 && this->dim0 > 0 && this->dim1 > 0);
-#endif
-    if (first_time)
-    {
-        int myprow=0, mypcol=0;
-        char order = 'R'; // row major process grid
-
-        blacs_ctxt = Csys2blacs_handle(comm_2D);
-        Cblacs_gridinit(&blacs_ctxt, &order, dim0, dim1);
-        Cblacs_gridinfo(blacs_ctxt, &dim0, &dim1, &myprow, &mypcol);
-    }
-    int ISRC = 0;
-    int info = 0;
-    descinit_(desc, &gr, &gc, &this->nb, &this->nb, &ISRC, &ISRC, &this->blacs_ctxt, &lld, &info);
-}
-
-int Parallel_2D::set_local2global(
-    const int& M_A,
-    const int& N_A,
-    std::ofstream& ofs_running,
-    std::ofstream& ofs_warning)
-{
-    ModuleBase::TITLE("Parallel_2D", "set_local2global");
-#ifdef __DEBUG
-    assert(M_A > 0 && N_A > 0);
-    assert(this->nb > 0);
-#endif
-
-    int dim[2];
-    int period[2];
-
-    // (0) every processor get it's id on the 2D comm
-    // : ( coord[0], coord[1] )
-    MPI_Cart_get(this->comm_2D, 2, dim, period, coord);
-    assert(dim[0] == this->dim0);
-    assert(dim[1] == this->dim1);
-
-    // local number of row and columns
+    // number of local rows and columns
     const int zero = 0;
-    nrow = numroc_(&M_A, &nb, &coord[0], &zero, &dim0);
-    ncol = numroc_(&N_A, &nb, &coord[1], &zero, &dim1);
-    nloc = nrow * ncol;
+    nrow = numroc_(&mg, &nb, &coord[0], &zero, &dim0);
+    ncol = numroc_(&ng, &nb, &coord[1], &zero, &dim1);
+    nloc = static_cast<int64_t>(nrow) * ncol;
 
-    // mohan add 2010-09-12
-    if (nrow == 0 || ncol == 0)
-    {
-        ofs_warning << " cpu 2D distribution : " << dim[0] << "*" << dim[1] << std::endl;
-        ofs_warning << " but, the number of row and column blocks are "
-                    << M_A / nb + (M_A % nb != 0) << "*" << N_A / nb + (N_A % nb != 0)
-                    << std::endl;
-        if (nb > 1)
-        {
-            return 1;
-        }
-        else
-        {
-            ModuleBase::WARNING_QUIT("Parallel_2D::set_local2global", "some processor has no blocks, try a smaller 'nb2d' parameter or reduce the number of mpi processes.");
-        }
-    }
+    // initialize the ScaLAPACK descriptor
+    int info = 0, lld = std::max(nrow, 1);
+    descinit_(desc, &mg, &ng, &nb, &nb, &zero, &zero, &blacs_ctxt, &lld, &info);
 
+    // generate the global-to-local and local-to-global index maps
     local2global_row_.resize(nrow);
+    global2local_row_ = std::vector<int>(mg, -1);
     for (int i = 0; i < nrow; ++i)
     {
-        local2global_row_[i] = (i / nb) * dim0 * nb + coord[0] * nb + i % nb;
+        local2global_row_[i] = (i / nb * dim0 + coord[0]) * nb + i % nb;
+        global2local_row_[local2global_row_[i]] = i;
     }
 
     local2global_col_.resize(ncol);
+    global2local_col_ = std::vector<int>(ng, -1);
     for (int j = 0; j < ncol; ++j)
     {
-        local2global_col_[j] = (j / nb) * dim1 * nb + coord[1] * nb + j % nb;
+        local2global_col_[j] = (j / nb * dim1 + coord[1]) * nb + j % nb;
+        global2local_col_[local2global_col_[j]] = j;
     }
+}
+
 
-    return 0;
+int Parallel_2D::init(
+    const int mg,
+    const int ng,
+    const int nb,
+    const MPI_Comm comm,
+    const bool mode
+)
+{
+    _init_proc_grid(comm, mode);
+    _set_dist_info(mg, ng, nb);
+    return nrow == 0 || ncol == 0;
 }
-#else
-void Parallel_2D::set_serial(const int& M_A, const int& N_A)
+
+
+int Parallel_2D::set(
+    const int mg,
+    const int ng,
+    const int nb,
+    const MPI_Comm comm_2D,
+    const int blacs_ctxt
+)
 {
-    ModuleBase::TITLE("Parallel_2D", "set_serial");
-    this->nrow = M_A;
-    this->ncol = N_A;
-    this->nloc = this->nrow * this->ncol;
-    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;
+    this->comm_2D = comm_2D;
+    this->blacs_ctxt = blacs_ctxt;
+    Cblacs_gridinfo(blacs_ctxt, &dim0, &dim1, &coord[0], &coord[1]);
+    _set_dist_info(mg, ng, nb);
+    return nrow == 0 || ncol == 0;
 }
 #endif
+
+void Parallel_2D::set_serial(const int mg, const int ng)
+{
+    assert(mg > 0 && ng > 0);
+
+    nb = 1;
+    dim0 = dim1 = 1;
+    coord[0] = coord[1] = 0;
+    nrow = mg;
+    ncol = ng;
+    nloc = nrow * ncol;
+    local2global_row_.resize(nrow);
+    local2global_col_.resize(ncol);
+    std::iota(local2global_row_.begin(), local2global_row_.end(), 0);
+    std::iota(local2global_col_.begin(), local2global_col_.end(), 0);
+    global2local_row_ = local2global_row_;
+    global2local_col_ = local2global_col_;
+#ifdef __MPI
+    comm_2D = MPI_COMM_NULL;
+    blacs_ctxt = -1;
+#endif
+}
+
diff --git a/source/module_basis/module_ao/parallel_2d.h b/source/module_basis/module_ao/parallel_2d.h
index 7ec1fa3a15..9905973698 100644
--- a/source/module_basis/module_ao/parallel_2d.h
+++ b/source/module_basis/module_ao/parallel_2d.h
@@ -1,7 +1,8 @@
 #ifndef _PARALLEL_2D_H_
 #define _PARALLEL_2D_H_
-#include "module_base/global_function.h"
-#include "module_base/global_variable.h"
+
+#include <vector>
+#include <cstdint>
 
 #ifdef __MPI
 #include <mpi.h>
@@ -12,117 +13,111 @@
 class Parallel_2D
 {
 public:
-    Parallel_2D();
-    ~Parallel_2D();
-
-    /// local size (nloc = nrow * ncol)
-    int nrow;
-    int ncol;
-    long nloc;
+    Parallel_2D() = default;
+    ~Parallel_2D() = default;
 
-    /// block size,
-    /// the default value of nb is 1,
-    /// but can change to larger value from input.
-    int nb = 1;
-
-    /// the number of processors in each dimension of MPI_Cart structure
-    int dim0;
-    int dim1;
-    /// the coordinate of current processor in each dimension of MPI_Cart structure
-    int coord[2];
+    /// number of local rows
+    int get_row_size() const { return nrow; };
 
-    /// test parameter
-    int testpb;
+    /// number of local columns
+    int get_col_size() const { return ncol; };
 
-    /// total number of columns of matrix in this processor
-    int get_col_size()const { return this->ncol; };
-
-    /// total number of rows of matrix in this processor
-    int get_row_size()const { return this->nrow; };
-
-    /// total number of elements of matrix in this processor
-    int get_local_size()const { return this->nloc; };
+    /// number of local matrix elements
+    int64_t get_local_size() const { return nloc; };
 
     /// get the local index of a global index (row)
-    int global2local_row(const int& igr) const
-    {
-        return this->global2local_row_[igr];
-    }
+    int global2local_row(const int igr) const { return global2local_row_[igr]; }
 
     /// get the local index of a global index (col)
-    int global2local_col(const int& igc) const
-    {
-        return this->global2local_col_[igc];
-    }
+    int global2local_col(const int igc) const { return global2local_col_[igc]; }
 
     /// get the global index of a local index (row)
-    int local2global_row(const int& ilr) const
-    {
-        return this->local2global_row_[ilr];
-    }
+    int local2global_row(const int ilr) const { return 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;
+    int local2global_col(const int ilc) const { return local2global_col_[ilc]; }
 
-    void set_block_size(const int& nb_in) { this->nb = nb_in; };
-    int get_block_size()const { return this->nb; };
+    /// check whether a global index is in this process
+    bool in_this_processor(const int iw1_all, const int iw2_all) const;
 
-    /// Set the 2D-structure of processors in each dimension.
-    /// dim0 and dim1 will be set as close to sqrt(nproc) as possible.
-    /// For example: nproc = 12,
-    /// if mode==0, d dim0 = 3, dim1 = 4; else, dim0 = 3, dim1 = 3.
-    void set_proc_dim(const int& dsize, bool mode = 0);
+    /// side length of 2d square block
+    int get_block_size() const { return nb; };
 
 #ifdef __MPI
-    int blacs_ctxt;    ///< blacs info
-    int desc[9];    ///<for matrix, nlocal*nlocal    
-    MPI_Comm comm_2D;   ///<communicator for 2D-block
-
-    /// create the 'comm_2D' stratege.
-    void mpi_create_cart(const MPI_Comm& diag_world);
-
-    /// set the map from local index to global index,
-    /// and set local sizes (nrow, ncol, nloc) by the way
-    int set_local2global(const int& M_A/**< global row size*/,
-        const int& N_A/**< global col size*/,
-        std::ofstream& ofs_running,
-        std::ofstream& ofs_warning);
-
-    ///@brief set the desc[9] of the 2D-block-cyclic distribution
-    void set_desc(const int& gr/**< global row size*/,
-        const int& gc/**< global col size*/,
-        const int& lld/**< leading local dimension*/,
-        bool first_time = true/**< true: call `Cblacs_get`; false: use `this->blacs_ctxt`*/);
-#else
-    void set_serial(const int& M_A/**< global row size*/,
-        const int& N_A/**< global col size*/);
+    /**
+     * @brief Initialize a BLACS grid with the given MPI communicator
+     * and set up the info of a block-cyclic distribution.
+     *
+     */
+    int init(
+        const int mg,
+        const int ng,
+        const int nb, // square block is assumed
+        const MPI_Comm comm,
+        bool mode = 0
+    );
+
+    /**
+     * @brief Set up the info of a block-cyclic distribution using given
+     * MPI communicator and BLACS context.
+     *
+     */
+    int set(
+        const int mg,
+        const int ng,
+        const int nb, // square block is assumed
+        const MPI_Comm comm_2D,
+        const int blacs_ctxt
+    );
+
+    /// BLACS context
+    int blacs_ctxt = -1;
+
+    /// ScaLAPACK descriptor
+    int desc[9] = {};
+
+    /// 2D Cartesian MPI communicator
+    MPI_Comm comm_2D = MPI_COMM_NULL;
 #endif
 
-    void set_global2local(const int& M_A,
-        const int& N_A,
-        const bool& div_2d,
-        std::ofstream& ofs_running);
+    void set_serial(const int mg, const int ng);
+
+    // FIXME the following variables should be private, but they are
+    // widely used in the code currently. Public visibility is kept
+    // for now, but might be changed in the future.
+
+    /// local size (nloc = nrow * ncol)
+    int nrow = 0;
+    int ncol = 0;
+    int64_t nloc = 0;
+
+    /// block size
+    int nb = 1;
+
+    /// number of processes in each dimension of the MPI Cartesian grid
+    int dim0 = 0;
+    int dim1 = 0;
+
+    /// process coordinate in the MPI Cartesian grid
+    int coord[2] = {-1, -1};
+
+    /// test parameter
+    int testpb = 0;
+
 
 protected:
 
     /// map from global index to local index
-    int* global2local_row_ = nullptr;
-    int* global2local_col_ = nullptr;
+    std::vector<int> global2local_row_;
+    std::vector<int> global2local_col_;
 
     /// 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_row_;
     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*/,
-        std::ofstream& ofs_running);
+#ifdef __MPI
+    void _init_proc_grid(const MPI_Comm comm, const bool mode);
+    void _set_dist_info(const int mg, const int ng, const int nb);
+#endif
 };
-#endif
\ No newline at end of file
+#endif
diff --git a/source/module_basis/module_ao/parallel_orbitals.cpp b/source/module_basis/module_ao/parallel_orbitals.cpp
index 67b055582b..ae45c1eea2 100644
--- a/source/module_basis/module_ao/parallel_orbitals.cpp
+++ b/source/module_basis/module_ao/parallel_orbitals.cpp
@@ -1,12 +1,8 @@
 #include "parallel_orbitals.h"
 
-#ifdef __MPI
-extern "C"
-{
 #include "module_base/blacs_connector.h"
 #include "module_base/scalapack_connector.h"
-}
-#endif
+#include "module_base/global_function.h"
 
 Parallel_Orbitals::Parallel_Orbitals()
 {
diff --git a/source/module_basis/module_ao/parallel_orbitals.h b/source/module_basis/module_ao/parallel_orbitals.h
index 15d0c7c7cf..f0d94de07e 100644
--- a/source/module_basis/module_ao/parallel_orbitals.h
+++ b/source/module_basis/module_ao/parallel_orbitals.h
@@ -1,6 +1,7 @@
 #ifndef _PARALLEL_ORBITALS_H_
 #define _PARALLEL_ORBITALS_H_
 #include "parallel_2d.h"
+#include <fstream>
 
 /// This class packs the information of 2D-block-cyclic for LCAO code:
 /// parallel distribution of basis, wavefunction and matrix.
@@ -95,4 +96,4 @@ class Parallel_Orbitals : public Parallel_2D
     const int* iat2iwt_ = nullptr;
 
 };
-#endif
\ No newline at end of file
+#endif
diff --git a/source/module_basis/module_ao/test/parallel_2d_test.cpp b/source/module_basis/module_ao/test/parallel_2d_test.cpp
index df386e6da5..8e73713e38 100644
--- a/source/module_basis/module_ao/test/parallel_2d_test.cpp
+++ b/source/module_basis/module_ao/test/parallel_2d_test.cpp
@@ -6,20 +6,9 @@
 
  /* Tested functions (in order):
   *
-  * - set_proc_dim
-  *   set the 2D-structure of processors in each dimension.
-  *
-  * - mpi_create_cart (parallel)
-  *   create a 2D MPI_Cart structure.
-  *
-  * - set_local2global (parallel)
-  *   set the map from local index to global index and the local sizes.
-  *
-  * - set_desc (parallel)
-  *   set the desc[9] of the 2D-block-cyclic distribution.
-  *
-  * - set_global2local
-  *   set the map from global index to local index (init, reuse).
+  * - init
+  *   initialize the MPI & BLACS 2d Cartesian grid and set up
+  *   the info of block-cyclic distribution.
   *
   * - set_serial (serial)
   *   set the local(=global) sizes.
@@ -40,19 +29,14 @@ class test_para2d : public testing::Test
     int my_rank = 0;
     std::vector<std::pair<int, int>> sizes{ {30, 35}, {49, 94}, {57, 57} };
     std::vector<int> nbs{ 1,2,3 };
-    std::ofstream ofs_running;
 #ifdef __MPI
     void SetUp() override
     {
         MPI_Comm_size(MPI_COMM_WORLD, &dsize);
         MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);
-        this->ofs_running.open("log" + std::to_string(my_rank) + ".txt");
-        ofs_running << "dsize(nproc) = " << dsize << std::endl;
-        ofs_running << "my_rank = " << my_rank << std::endl;
     }
     void TearDown() override
     {
-        ofs_running.close();
     }
 #endif
 };
@@ -67,23 +51,22 @@ TEST_F(test_para2d, Divide2D)
         for (auto nb : nbs)
         {
             Parallel_2D p2d;
-            p2d.set_block_size(nb);
-            EXPECT_EQ(p2d.get_block_size(), nb);
 
             for (auto mode : { 0,1 })
             {
-                //1. set dim0 and dim1
-                p2d.set_proc_dim(dsize, mode);
+                p2d.init(gr, gc, nb, MPI_COMM_WORLD, mode);
+
+                EXPECT_EQ(p2d.get_block_size(), nb);
+
+                //1. dim0 and dim1
                 EXPECT_EQ(p2d.dim0 * p2d.dim1, dsize);
                 if (mode)EXPECT_LE(p2d.dim1, p2d.dim0);
                 else EXPECT_LE(p2d.dim0, p2d.dim1);
 
-                //2. mpi_create_cart
-                p2d.mpi_create_cart(MPI_COMM_WORLD);
+                //2. MPI 2d communicator
                 EXPECT_NE(p2d.comm_2D, MPI_COMM_NULL);
 
-                //3. set_local2global and local sizes
-                p2d.set_local2global(gr, gc, ofs_running, ofs_running);
+                //3. local2global and local sizes
                 int lr = p2d.get_row_size();
                 int lc = p2d.get_col_size();
                 EXPECT_EQ(lr * lc, p2d.get_local_size());
@@ -96,8 +79,7 @@ TEST_F(test_para2d, Divide2D)
                 EXPECT_EQ(lr, cal_lsize(gr, nb, p2d.dim0, p2d.coord[0]));
                 EXPECT_EQ(lc, cal_lsize(gc, nb, p2d.dim1, p2d.coord[1]));
 
-                //4. set_desc
-                p2d.set_desc(gr, gc, lr);
+                //4. ScaLAPACK descriptor
                 EXPECT_EQ(p2d.desc[0], 1);
                 EXPECT_EQ(p2d.desc[1], p2d.blacs_ctxt);
                 EXPECT_EQ(p2d.desc[2], gr);
@@ -108,8 +90,7 @@ TEST_F(test_para2d, Divide2D)
                 EXPECT_EQ(p2d.desc[7], 0);
                 EXPECT_EQ(p2d.desc[8], lr);
 
-                //5. set_global2local
-                p2d.set_global2local(gr, gc, true, ofs_running);
+                //5. global2local
                 auto sum_array = [&p2d](const int& gr, const int& gc) -> std::pair<int, int>
                     {
                         int sum_row = 0; int sum_col = 0;
@@ -135,30 +116,15 @@ TEST_F(test_para2d, DescReuseCtxt)
     for (auto nb : nbs)
     {
         Parallel_2D p1;
-        p1.set_block_size(nb);
-        p1.set_proc_dim(dsize);
-        p1.mpi_create_cart(MPI_COMM_WORLD);
-        p1.set_local2global(sizes[0].first, sizes[0].second, ofs_running, ofs_running);
-        p1.set_desc(sizes[0].first, sizes[0].second, p1.get_row_size());
+        p1.init(sizes[0].first, sizes[0].second, nb, MPI_COMM_WORLD);
 
         Parallel_2D p2;    // use 2 different sizes, but they can share the same ctxt
-        p2.set_block_size(nb);
-        p2.set_proc_dim(dsize);
-        p2.comm_2D = p1.comm_2D;
-        p2.blacs_ctxt = p1.blacs_ctxt;
-        p2.set_local2global(sizes[1].first, sizes[1].second, ofs_running, ofs_running);
-        p2.set_desc(sizes[1].first, sizes[1].second, p2.get_row_size(), false);
+        p2.set(sizes[1].first, sizes[1].second, nb, p1.comm_2D, p1.blacs_ctxt);
 
         EXPECT_EQ(p1.desc[1], p2.desc[1]);
 
-        Parallel_2D p3;    // using default `set_desc`, p3 can't share the same ctxt with p1
-        p3.set_block_size(nb);
-        p3.set_proc_dim(dsize);
-        p3.comm_2D = p1.comm_2D;
-        p3.blacs_ctxt = p1.blacs_ctxt;
-        p3.set_local2global(sizes[2].first, sizes[2].second, ofs_running, ofs_running);
-        p3.set_desc(sizes[2].first, sizes[2].second, p3.get_row_size());
-
+        Parallel_2D p3;
+        p3.init(sizes[2].first, sizes[2].second, nb, MPI_COMM_WORLD);
         EXPECT_NE(p1.desc[1], p3.desc[1]);
     }
 }
@@ -172,14 +138,12 @@ TEST_F(test_para2d, Serial)
 
         Parallel_2D p2d;
 
-        //1. set dim0 and dim1
-        p2d.set_proc_dim(1);
+        // set_serial
+        p2d.set_serial(gr, gc);
         EXPECT_EQ(p2d.dim0 * p2d.dim1, 1);
         EXPECT_EQ(p2d.dim0, 1);
         EXPECT_EQ(p2d.dim1, 1);
 
-        //2. set_serial
-        p2d.set_serial(gr, gc);
         EXPECT_EQ(p2d.get_row_size(), gr);
         EXPECT_EQ(p2d.get_col_size(), gc);
         EXPECT_EQ(p2d.get_local_size(), gr * gc);
@@ -188,8 +152,7 @@ TEST_F(test_para2d, Serial)
         for (int i = 0; i < gc; ++i)
             EXPECT_EQ(p2d.local2global_col(i), i);
 
-        //3. set_global2local
-        p2d.set_global2local(gr, gc, false, ofs_running);
+        //3. global2local
         for (int i = 0;i < gr;++i)
             EXPECT_EQ(p2d.global2local_row(i), i);
         for (int i = 0;i < gc;++i)
diff --git a/source/module_basis/module_ao/test/parallel_orbitals_test.cpp b/source/module_basis/module_ao/test/parallel_orbitals_test.cpp
index a51eddeeeb..06bd6f010b 100644
--- a/source/module_basis/module_ao/test/parallel_orbitals_test.cpp
+++ b/source/module_basis/module_ao/test/parallel_orbitals_test.cpp
@@ -26,19 +26,14 @@ class TestParaO : public testing::Test
     std::vector<std::pair<int, int>> sizes{ {50, 50} , {60, 60}};
     std::vector<int> nat{ 10, 5};
     std::vector<int> nbs{ 1,2,3 };
-    std::ofstream ofs_running;
 #ifdef __MPI
     void SetUp() override
     {
         MPI_Comm_size(MPI_COMM_WORLD, &dsize);
         MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);
-        this->ofs_running.open("log" + std::to_string(my_rank) + ".txt");
-        ofs_running << "dsize(nproc) = " << dsize << std::endl;
-        ofs_running << "my_rank = " << my_rank << std::endl;
     }
     void TearDown() override
     {
-        ofs_running.close();
     }
 #endif
 };
@@ -53,23 +48,21 @@ TEST_F(TestParaO, Divide2D)
         for (auto nb : nbs)
         {
             Parallel_Orbitals po;
-            po.set_block_size(nb);
-            EXPECT_EQ(po.get_block_size(), nb);
 
             for (auto mode : { 0,1 })
             {
-                //1. set dim0 and dim1
-                po.set_proc_dim(dsize, mode);
+                po.init(gr, gc, nb, MPI_COMM_WORLD, mode);
+                EXPECT_EQ(po.get_block_size(), nb);
+
+                //1. dim0 and dim1
                 EXPECT_EQ(po.dim0 * po.dim1, dsize);
                 if (mode)EXPECT_LE(po.dim1, po.dim0);
                 else EXPECT_LE(po.dim0, po.dim1);
 
-                //2. mpi_create_cart
-                po.mpi_create_cart(MPI_COMM_WORLD);
+                //2. comm_2D
                 EXPECT_NE(po.comm_2D, MPI_COMM_NULL);
 
-                //3. set_local2global and local sizes
-                po.set_local2global(gr, gc, ofs_running, ofs_running);
+                //3. local2global and local sizes
                 int lr = po.get_row_size();
                 int lc = po.get_col_size();
                 EXPECT_EQ(lr * lc, po.get_local_size());
@@ -82,8 +75,7 @@ TEST_F(TestParaO, Divide2D)
                 EXPECT_EQ(lr, cal_lsize(gr, nb, po.dim0, po.coord[0]));
                 EXPECT_EQ(lc, cal_lsize(gc, nb, po.dim1, po.coord[1]));
 
-                //4. set_desc
-                po.set_desc(gr, gc, lr);
+                //4. ScaLAPACK descriptor
                 EXPECT_EQ(po.desc[0], 1);
                 EXPECT_EQ(po.desc[1], po.blacs_ctxt);
                 EXPECT_EQ(po.desc[2], gr);
@@ -94,8 +86,7 @@ TEST_F(TestParaO, Divide2D)
                 EXPECT_EQ(po.desc[7], 0);
                 EXPECT_EQ(po.desc[8], lr);
 
-                //5. set_global2local
-                po.set_global2local(gr, gc, true, ofs_running);
+                //5. global2local
                 auto sum_array = [&po](const int& gr, const int& gc) -> std::pair<int, int>
                 {
                     int sum_row = 0; int sum_col = 0;
@@ -168,17 +159,14 @@ TEST_F(TestParaO, Serial)
         Parallel_Orbitals po;
 
         //1. set dim0 and dim1
-        po.set_proc_dim(1);
-        EXPECT_EQ(po.dim0 * po.dim1, 1);
-
         //2. set_serial
         po.set_serial(gr, gc);
+        EXPECT_EQ(po.dim0 * po.dim1, 1);
         EXPECT_EQ(po.get_row_size(), gr);
         EXPECT_EQ(po.get_col_size(), gc);
         EXPECT_EQ(po.get_local_size(), gr * gc);
 
-        //3. set_global2local
-        po.set_global2local(gr, gc, false, ofs_running);
+        //3. global2local
         for (int i = 0;i < gr;++i)
             EXPECT_EQ(po.global2local_row(i), i);
         for (int i = 0;i < gc;++i)
diff --git a/source/module_elecstate/module_dm/test/test_cal_dm_R.cpp b/source/module_elecstate/module_dm/test/test_cal_dm_R.cpp
index b3a0f71fc6..3a4f68e74a 100644
--- a/source/module_elecstate/module_dm/test/test_cal_dm_R.cpp
+++ b/source/module_elecstate/module_dm/test/test_cal_dm_R.cpp
@@ -81,18 +81,12 @@ class DMTest : public testing::Test
 #ifdef __MPI
     void init_parav()
     {
+        int nb = 2;
         int global_row = test_size * test_nw;
         int global_col = test_size * test_nw;
         std::ofstream ofs_running;
         paraV = new Parallel_Orbitals();
-        paraV->set_block_size(2 /* nb_2d set to be 2*/);
-        paraV->set_proc_dim(dsize, 0);
-        paraV->mpi_create_cart(MPI_COMM_WORLD);
-        paraV->set_local2global(global_row, global_col, ofs_running, ofs_running);
-        int lr = paraV->get_row_size();
-        int lc = paraV->get_col_size();
-        paraV->set_desc(global_row, global_col, lr);
-        paraV->set_global2local(global_row, global_col, true, ofs_running);
+        paraV->init(global_row, global_col, nb, MPI_COMM_WORLD);
         paraV->set_atomic_trace(ucell.get_iat2iwt(), test_size, global_row);
     }
 #else
@@ -341,4 +335,4 @@ int main(int argc, char** argv)
     MPI_Finalize();
 #endif
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_elecstate/module_dm/test/test_cal_dmk_psi.cpp b/source/module_elecstate/module_dm/test/test_cal_dmk_psi.cpp
index 6697ce2f63..16c946200d 100644
--- a/source/module_elecstate/module_dm/test/test_cal_dmk_psi.cpp
+++ b/source/module_elecstate/module_dm/test/test_cal_dmk_psi.cpp
@@ -80,18 +80,12 @@ class DMTest : public testing::Test
 #ifdef __MPI
     void init_parav()
     {
+        int nb = 2;
         int global_row = test_size * test_nw;
         int global_col = test_size * test_nw;
         std::ofstream ofs_running;
         paraV = new Parallel_Orbitals();
-        paraV->set_block_size(2 /* nb_2d set to be 2*/);
-        paraV->set_proc_dim(dsize, 0);
-        paraV->mpi_create_cart(MPI_COMM_WORLD);
-        paraV->set_local2global(global_row, global_col, ofs_running, ofs_running);
-        int lr = paraV->get_row_size();
-        int lc = paraV->get_col_size();
-        paraV->set_desc(global_row, global_col, lr);
-        paraV->set_global2local(global_row, global_col, true, ofs_running);
+        paraV->init(global_row, global_col, nb, MPI_COMM_WORLD);
         paraV->set_atomic_trace(ucell.get_iat2iwt(), test_size, global_row);
     }
 #else
@@ -180,4 +174,4 @@ int main(int argc, char** argv)
     MPI_Finalize();
 #endif
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_elecstate/module_dm/test/test_dm_R_init.cpp b/source/module_elecstate/module_dm/test/test_dm_R_init.cpp
index e3f1441898..2a8477d217 100644
--- a/source/module_elecstate/module_dm/test/test_dm_R_init.cpp
+++ b/source/module_elecstate/module_dm/test/test_dm_R_init.cpp
@@ -81,18 +81,12 @@ class DMTest : public testing::Test
 #ifdef __MPI
     void init_parav()
     {
+        int nb = 2;
         int global_row = test_size * test_nw;
         int global_col = test_size * test_nw;
         std::ofstream ofs_running;
         paraV = new Parallel_Orbitals();
-        paraV->set_block_size(2 /* nb_2d set to be 2*/);
-        paraV->set_proc_dim(dsize, 0);
-        paraV->mpi_create_cart(MPI_COMM_WORLD);
-        paraV->set_local2global(global_row, global_col, ofs_running, ofs_running);
-        int lr = paraV->get_row_size();
-        int lc = paraV->get_col_size();
-        paraV->set_desc(global_row, global_col, lr);
-        paraV->set_global2local(global_row, global_col, true, ofs_running);
+        paraV->init(global_row, global_col, nb, MPI_COMM_WORLD);
         paraV->set_atomic_trace(ucell.get_iat2iwt(), test_size, global_row);
     }
 #else
@@ -321,4 +315,4 @@ int main(int argc, char** argv)
     MPI_Finalize();
 #endif
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_elecstate/module_dm/test/test_dm_constructor.cpp b/source/module_elecstate/module_dm/test/test_dm_constructor.cpp
index 89466594d4..5482a5687e 100644
--- a/source/module_elecstate/module_dm/test/test_dm_constructor.cpp
+++ b/source/module_elecstate/module_dm/test/test_dm_constructor.cpp
@@ -80,18 +80,12 @@ class DMTest : public testing::Test
 #ifdef __MPI
     void init_parav()
     {
+        int nb = 2;
         int global_row = test_size * test_nw;
         int global_col = test_size * test_nw;
         std::ofstream ofs_running;
         paraV = new Parallel_Orbitals();
-        paraV->set_block_size(2 /* nb_2d set to be 2*/);
-        paraV->set_proc_dim(dsize, 0);
-        paraV->mpi_create_cart(MPI_COMM_WORLD);
-        paraV->set_local2global(global_row, global_col, ofs_running, ofs_running);
-        int lr = paraV->get_row_size();
-        int lc = paraV->get_col_size();
-        paraV->set_desc(global_row, global_col, lr);
-        paraV->set_global2local(global_row, global_col, true, ofs_running);
+        paraV->init(global_row, global_col, nb, MPI_COMM_WORLD);
         paraV->set_atomic_trace(ucell.get_iat2iwt(), test_size, global_row);
     }
 #else
@@ -260,4 +254,4 @@ int main(int argc, char** argv)
     MPI_Finalize();
 #endif
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_elecstate/module_dm/test/test_dm_io.cpp b/source/module_elecstate/module_dm/test/test_dm_io.cpp
index 0633c1c2ee..e264b156a6 100644
--- a/source/module_elecstate/module_dm/test/test_dm_io.cpp
+++ b/source/module_elecstate/module_dm/test/test_dm_io.cpp
@@ -123,18 +123,12 @@ class DMTest : public testing::Test
 #ifdef __MPI
     void init_parav()
     {
+        int nb = 2;
         int global_row = test_size * test_nw;
         int global_col = test_size * test_nw;
         std::ofstream ofs_running;
         paraV = new Parallel_Orbitals();
-        paraV->set_block_size(2 /* nb_2d set to be 2*/);
-        paraV->set_proc_dim(dsize, 0);
-        paraV->mpi_create_cart(MPI_COMM_WORLD);
-        paraV->set_local2global(global_row, global_col, ofs_running, ofs_running);
-        int lr = paraV->get_row_size();
-        int lc = paraV->get_col_size();
-        paraV->set_desc(global_row, global_col, lr);
-        paraV->set_global2local(global_row, global_col, true, ofs_running);
+        paraV->init(global_row, global_col, nb, MPI_COMM_WORLD);
         paraV->set_atomic_trace(ucell->get_iat2iwt(), test_size, global_row);
     }
 #else
@@ -195,4 +189,4 @@ int main(int argc, char** argv)
     MPI_Finalize();
 #endif
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_esolver/io_npz.cpp b/source/module_esolver/io_npz.cpp
index 3415ae3004..8a557884fe 100644
--- a/source/module_esolver/io_npz.cpp
+++ b/source/module_esolver/io_npz.cpp
@@ -45,15 +45,6 @@ void ESolver_KS_LCAO<TK, TR>::read_mat_npz(std::string& zipname, hamilt::HContai
         return;
     }
 
-/*
-    hamilt::HContainer<double>* HR_serial;
-    Parallel_Orbitals serialV;
-    serialV.set_proc_dim(1,0);
-    serialV.mpi_create_cart(MPI_COMM_WORLD);
-    serialV.set_local2global(GlobalV::NLOCAL, GlobalV::NLOCAL, GlobalV::ofs_running, GlobalV::ofs_warning);
-    serialV.set_global2local(GlobalV::NLOCAL, GlobalV::NLOCAL, false, GlobalV::ofs_running);
-    serialV.set_atomic_trace(GlobalC::ucell.get_iat2iwt(), GlobalC::ucell.nat, GlobalV::NLOCAL);
-*/
     if(GlobalV::MY_RANK == 0)
     {
         //HR_serial = new hamilt::HContainer<double>(&serialV);
@@ -435,7 +426,7 @@ void ESolver_KS_LCAO<TK, TR>::output_mat_npz(std::string& zipname, const hamilt:
 #ifdef __MPI
     hamilt::HContainer<double>* HR_serial;
     Parallel_Orbitals serialV;
-    serialV.set_global2local(GlobalV::NLOCAL, GlobalV::NLOCAL, false, GlobalV::ofs_running);
+    serialV.set_serial(GlobalV::NLOCAL, GlobalV::NLOCAL);
     serialV.set_atomic_trace(GlobalC::ucell.get_iat2iwt(), GlobalC::ucell.nat, GlobalV::NLOCAL);
     if(GlobalV::MY_RANK == 0)
     {
@@ -496,4 +487,4 @@ void ESolver_KS_LCAO<TK, TR>::output_mat_npz(std::string& zipname, const hamilt:
 template class ESolver_KS_LCAO<double, double>;
 template class ESolver_KS_LCAO<std::complex<double>, double>;
 template class ESolver_KS_LCAO<std::complex<double>, std::complex<double>>;
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_T_NL_cd.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_T_NL_cd.cpp
index 791083dc0c..cc30c21887 100644
--- a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_T_NL_cd.cpp
+++ b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_T_NL_cd.cpp
@@ -108,18 +108,12 @@ class TNLTest : public ::testing::Test
 #ifdef __MPI
     void init_parav()
     {
+        int nb = 20;
         int global_row = test_size * test_nw * 2;
         int global_col = test_size * test_nw * 2;
         std::ofstream ofs_running;
         paraV = new Parallel_Orbitals();
-        paraV->set_block_size(20/* nb_2d set to be 2*/);
-        paraV->set_proc_dim(dsize, 0);
-        paraV->mpi_create_cart(MPI_COMM_WORLD);
-        paraV->set_local2global(global_row, global_col, ofs_running, ofs_running);
-        int lr = paraV->get_row_size();
-        int lc = paraV->get_col_size();
-        paraV->set_desc(global_row, global_col, lr);
-        paraV->set_global2local(global_row, global_col, true, ofs_running);
+        paraV->init(global_row, global_col, nb, MPI_COMM_WORLD);
         paraV->set_atomic_trace(ucell.get_iat2iwt(), test_size, global_row);
     }
 #else
@@ -266,4 +260,4 @@ int main(int argc, char** argv)
     MPI_Finalize();
 #endif
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_dftu.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_dftu.cpp
index a62c409b47..92ebc20fad 100644
--- a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_dftu.cpp
+++ b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_dftu.cpp
@@ -117,18 +117,12 @@ class DFTUTest : public ::testing::Test
 #ifdef __MPI
     void init_parav()
     {
+        int nb = 10;
         int global_row = test_size * test_nw;
         int global_col = test_size * test_nw;
         std::ofstream ofs_running;
         paraV = new Parallel_Orbitals();
-        paraV->set_block_size(10/* nb_2d set to be 2*/);
-        paraV->set_proc_dim(dsize, 0);
-        paraV->mpi_create_cart(MPI_COMM_WORLD);
-        paraV->set_local2global(global_row, global_col, ofs_running, ofs_running);
-        int lr = paraV->get_row_size();
-        int lc = paraV->get_col_size();
-        paraV->set_desc(global_row, global_col, lr);
-        paraV->set_global2local(global_row, global_col, true, ofs_running);
+        paraV->init(global_row, global_col, nb, MPI_COMM_WORLD);
         paraV->set_atomic_trace(ucell.get_iat2iwt(), test_size, global_row);
     }
 #else
@@ -296,4 +290,4 @@ int main(int argc, char** argv)
     MPI_Finalize();
 #endif
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_ekineticnew.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_ekineticnew.cpp
index 5e34ac3417..4b8e13a22c 100644
--- a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_ekineticnew.cpp
+++ b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_ekineticnew.cpp
@@ -77,18 +77,12 @@ class EkineticNewTest : public ::testing::Test
 #ifdef __MPI
     void init_parav()
     {
+        int nb = 10;
         int global_row = test_size * test_nw;
         int global_col = test_size * test_nw;
         std::ofstream ofs_running;
         paraV = new Parallel_Orbitals();
-        paraV->set_block_size(10/* nb_2d set to be 2*/);
-        paraV->set_proc_dim(dsize, 0);
-        paraV->mpi_create_cart(MPI_COMM_WORLD);
-        paraV->set_local2global(global_row, global_col, ofs_running, ofs_running);
-        int lr = paraV->get_row_size();
-        int lc = paraV->get_col_size();
-        paraV->set_desc(global_row, global_col, lr);
-        paraV->set_global2local(global_row, global_col, true, ofs_running);
+        paraV->init(global_row, global_col, nb, MPI_COMM_WORLD);
         paraV->set_atomic_trace(ucell.get_iat2iwt(), test_size, global_row);
     }
 #else
@@ -219,4 +213,4 @@ int main(int argc, char** argv)
     MPI_Finalize();
 #endif
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_nonlocalnew.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_nonlocalnew.cpp
index b433cb86df..725c9d5008 100644
--- a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_nonlocalnew.cpp
+++ b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_nonlocalnew.cpp
@@ -106,18 +106,12 @@ class NonlocalNewTest : public ::testing::Test
 #ifdef __MPI
     void init_parav()
     {
+        int nb = 10;
         int global_row = test_size * test_nw;
         int global_col = test_size * test_nw;
         std::ofstream ofs_running;
         paraV = new Parallel_Orbitals();
-        paraV->set_block_size(10/* nb_2d set to be 2*/);
-        paraV->set_proc_dim(dsize, 0);
-        paraV->mpi_create_cart(MPI_COMM_WORLD);
-        paraV->set_local2global(global_row, global_col, ofs_running, ofs_running);
-        int lr = paraV->get_row_size();
-        int lc = paraV->get_col_size();
-        paraV->set_desc(global_row, global_col, lr);
-        paraV->set_global2local(global_row, global_col, true, ofs_running);
+        paraV->init(global_row, global_col, nb, MPI_COMM_WORLD);
         paraV->set_atomic_trace(ucell.get_iat2iwt(), test_size, global_row);
     }
 #else
@@ -265,4 +259,4 @@ int main(int argc, char** argv)
     MPI_Finalize();
 #endif
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_overlapnew.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_overlapnew.cpp
index 423d4b727b..4350d4ba10 100644
--- a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_overlapnew.cpp
+++ b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_overlapnew.cpp
@@ -77,18 +77,12 @@ class OverlapNewTest : public ::testing::Test
 #ifdef __MPI
     void init_parav()
     {
+        int nb = 10;
         int global_row = test_size * test_nw;
         int global_col = test_size * test_nw;
         std::ofstream ofs_running;
         paraV = new Parallel_Orbitals();
-        paraV->set_block_size(10/* nb_2d set to be 2*/);
-        paraV->set_proc_dim(dsize, 0);
-        paraV->mpi_create_cart(MPI_COMM_WORLD);
-        paraV->set_local2global(global_row, global_col, ofs_running, ofs_running);
-        int lr = paraV->get_row_size();
-        int lc = paraV->get_col_size();
-        paraV->set_desc(global_row, global_col, lr);
-        paraV->set_global2local(global_row, global_col, true, ofs_running);
+        paraV->init(global_row, global_col, nb, MPI_COMM_WORLD);
         paraV->set_atomic_trace(ucell.get_iat2iwt(), test_size, global_row);
     }
 #else
@@ -207,4 +201,4 @@ int main(int argc, char** argv)
     MPI_Finalize();
 #endif
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_overlapnew_cd.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_overlapnew_cd.cpp
index 7c2af7cf87..1c0187c850 100644
--- a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_overlapnew_cd.cpp
+++ b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_overlapnew_cd.cpp
@@ -76,18 +76,12 @@ class OverlapNewTest : public ::testing::Test
 #ifdef __MPI
     void init_parav()
     {
+        int nb = 10;
         int global_row = test_size * test_nw * 2;
         int global_col = test_size * test_nw * 2;
         std::ofstream ofs_running;
         paraV = new Parallel_Orbitals();
-        paraV->set_block_size(10/* nb_2d set to be 2*/);
-        paraV->set_proc_dim(dsize, 0);
-        paraV->mpi_create_cart(MPI_COMM_WORLD);
-        paraV->set_local2global(global_row, global_col, ofs_running, ofs_running);
-        int lr = paraV->get_row_size();
-        int lc = paraV->get_col_size();
-        paraV->set_desc(global_row, global_col, lr);
-        paraV->set_global2local(global_row, global_col, true, ofs_running);
+        paraV->init(global_row, global_col, nb, MPI_COMM_WORLD);
         paraV->set_atomic_trace(ucell.get_iat2iwt(), test_size, global_row);
     }
 #else
@@ -205,4 +199,4 @@ int main(int argc, char** argv)
     MPI_Finalize();
 #endif
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_sc_lambda_lcao.cpp b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_sc_lambda_lcao.cpp
index 2290b9486b..15efa6c7b8 100644
--- a/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_sc_lambda_lcao.cpp
+++ b/source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/test/test_sc_lambda_lcao.cpp
@@ -84,7 +84,7 @@ TEST_F(ScLambdaLCAOTest, ContributeHk)
     std::ofstream ofs("test.log");
     int nrow = 2;
     int ncol = 2;
-    paraV.set_global2local(nrow, ncol, false, ofs);
+    paraV.set_serial(nrow, ncol);
     EXPECT_EQ(paraV.nloc, 4);
     // set LM
     LCAO_Matrix LM;
@@ -151,7 +151,7 @@ TEST_F(ScLambdaLCAOTest, ContributeHkS2)
     std::ofstream ofs("test.log");
     int nrow = 1;
     int ncol = 1;
-    paraV.set_global2local(nrow, ncol, false, ofs);
+    paraV.set_serial(nrow, ncol);
     EXPECT_EQ(paraV.nloc, 1);
     // set LM
     LCAO_Matrix LM;
@@ -218,4 +218,4 @@ TEST_F(ScLambdaLCAOTest, TemplateHelpers)
     EXPECT_NO_THROW(sc_lambda_op1.contributeHR());
     EXPECT_NO_THROW(sc_lambda_op2.contributeHR());
     EXPECT_NO_THROW(sc_lambda_op2.contributeHk(0));
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.h b/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.h
index 7b501e790a..3516937597 100644
--- a/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.h
+++ b/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test.h
@@ -9,7 +9,7 @@
 #include "module_cell/module_neighbor/sltk_grid_driver.h"
 #include "module_cell/module_neighbor/sltk_atom_arrange.h"
 #include "klist.h"
-#include "parallel_orbitals.h"
+//#include "parallel_orbitals.h"
 
 #include "../LCAO_deepks.h"
 
diff --git a/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test_prep.cpp b/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test_prep.cpp
index 9135c66523..13fa791391 100644
--- a/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test_prep.cpp
+++ b/source/module_hamilt_lcao/module_deepks/test/LCAO_deepks_test_prep.cpp
@@ -13,8 +13,7 @@ void test_deepks::preparation()
 	this->set_orbs(ucell.lat0);
 	this->prep_neighbour();
 
-    this->ParaO.set_global2local(GlobalV::NLOCAL, GlobalV::NLOCAL,
-        false, GlobalV::ofs_running);
+    this->ParaO.set_serial(GlobalV::NLOCAL, GlobalV::NLOCAL);
 }
 
 void test_deepks::set_parameters()
diff --git a/source/module_hamilt_lcao/module_deltaspin/test/cal_h_lambda_test.cpp b/source/module_hamilt_lcao/module_deltaspin/test/cal_h_lambda_test.cpp
index f2e488e893..ba7d17527d 100644
--- a/source/module_hamilt_lcao/module_deltaspin/test/cal_h_lambda_test.cpp
+++ b/source/module_hamilt_lcao/module_deltaspin/test/cal_h_lambda_test.cpp
@@ -29,7 +29,7 @@ TEST_F(SpinConstrainTest, CalHLambda)
     int nrow = 2;
     int ncol = 2;
     sc.set_nspin(4);
-    paraV.set_global2local(nrow, ncol, false, ofs);
+    paraV.set_serial(nrow, ncol);
     sc.set_ParaV(&paraV);
     EXPECT_EQ(sc.ParaV->nloc, 4);
     std::map<int, int> atomCounts = {
@@ -98,7 +98,7 @@ TEST_F(SpinConstrainTest, CalHLambdaS2)
     int nrow = 1;
     int ncol = 1;
     sc.set_nspin(2);
-    paraV.set_global2local(nrow, ncol, false, ofs);
+    paraV.set_serial(nrow, ncol);
     sc.set_ParaV(&paraV);
     EXPECT_EQ(sc.ParaV->nloc, 1);
     std::map<int, int> atomCounts = {
@@ -152,4 +152,4 @@ TEST_F(SpinConstrainTest, CalHLambdaS2)
     EXPECT_DOUBLE_EQ(h_lambda[0].real(), rowMajor_h_lambda1[0].real());
     EXPECT_DOUBLE_EQ(h_lambda[0].imag(), rowMajor_h_lambda1[0].imag());
     remove("test.log");
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/module_deltaspin/test/cal_mw_helper_test.cpp b/source/module_hamilt_lcao/module_deltaspin/test/cal_mw_helper_test.cpp
index 79b7ec70cd..ceab7ae388 100644
--- a/source/module_hamilt_lcao/module_deltaspin/test/cal_mw_helper_test.cpp
+++ b/source/module_hamilt_lcao/module_deltaspin/test/cal_mw_helper_test.cpp
@@ -77,7 +77,7 @@ TEST_F(SpinConstrainTest, CollectMW)
     int nrow = 2;
     int ncol = 2;
     std::ofstream ofs("test.log");
-    paraV.set_global2local(nrow, ncol, false, ofs);
+    paraV.set_serial(nrow, ncol);
     sc.set_ParaV(&paraV);
     // Prepare the input data
     int nw = sc.get_nw();
@@ -153,7 +153,7 @@ TEST_F(SpinConstrainTest, CollectMWS2)
     int nrow = 1;
     int ncol = 1;
     std::ofstream ofs("test.log");
-    paraV.set_global2local(nrow, ncol, false, ofs);
+    paraV.set_serial(nrow, ncol);
     sc.set_ParaV(&paraV);
     // Prepare the input data
     int nw = sc.get_nw();
@@ -178,4 +178,4 @@ TEST_F(SpinConstrainTest, CollectMWS2)
         }
     }
     remove("test.log");
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/module_deltaspin/test/spin_constrain_test.cpp b/source/module_hamilt_lcao/module_deltaspin/test/spin_constrain_test.cpp
index b05ae34a82..ef514c24c6 100644
--- a/source/module_hamilt_lcao/module_deltaspin/test/spin_constrain_test.cpp
+++ b/source/module_hamilt_lcao/module_deltaspin/test/spin_constrain_test.cpp
@@ -386,7 +386,7 @@ TYPED_TEST(SpinConstrainTest, SetParaV)
     int nrow = 4;
     int ncol = 4;
     std::ofstream ofs("test.log");
-    paraV.set_global2local(nrow, ncol, false, ofs);
+    paraV.set_serial(nrow, ncol);
     this->sc.set_ParaV(&paraV);
     EXPECT_EQ(this->sc.ParaV->nloc, nrow * ncol);
     remove("test.log");
@@ -406,4 +406,4 @@ TYPED_TEST(SpinConstrainTest, PrintMi)
     output = testing::internal::GetCapturedStdout();
     EXPECT_THAT(output, testing::HasSubstr("Total Magnetism (uB):"));
     EXPECT_THAT(output, testing::HasSubstr("ATOM      0         0.0000000000"));
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/module_hcontainer/atom_pair.cpp b/source/module_hamilt_lcao/module_hcontainer/atom_pair.cpp
index d308bfd8e9..cd0ead0435 100644
--- a/source/module_hamilt_lcao/module_hcontainer/atom_pair.cpp
+++ b/source/module_hamilt_lcao/module_hcontainer/atom_pair.cpp
@@ -1,5 +1,6 @@
 #include "atom_pair.h"
 #include <complex>
+#include <cassert>
 #include "module_base/blas_connector.h"
 
 namespace hamilt
@@ -797,4 +798,4 @@ size_t AtomPair<T>::get_memory_size() const
 template class AtomPair<double>;
 template class AtomPair<std::complex<double>>;
 
-} // namespace hamilt
\ No newline at end of file
+} // namespace hamilt
diff --git a/source/module_hamilt_lcao/module_hcontainer/atom_pair.h b/source/module_hamilt_lcao/module_hcontainer/atom_pair.h
index 309231d79b..d5f7d69af9 100644
--- a/source/module_hamilt_lcao/module_hcontainer/atom_pair.h
+++ b/source/module_hamilt_lcao/module_hcontainer/atom_pair.h
@@ -6,6 +6,8 @@
 #include "module_basis/module_ao/parallel_orbitals.h"
 
 #include <vector>
+#include <complex>
+#include <tuple>
 
 namespace hamilt
 {
@@ -300,4 +302,4 @@ class AtomPair
 
 } // namespace hamilt
 
-#endif
\ No newline at end of file
+#endif
diff --git a/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer.cpp b/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer.cpp
index 2624182556..8473451dcc 100644
--- a/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer.cpp
+++ b/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer.cpp
@@ -572,7 +572,7 @@ TEST_F(HContainerTest, atompair_funcs)
     // get_matrix_value will use global2local_row and global2local_col in Parallel_Orbitals
     // so we need to set them
     std::ofstream ofs("test_hcontainer.log");
-    PO.set_global2local(4, 4, false, ofs);
+    PO.set_serial(4, 4);
     // joint 4 2*2 BaseMatrix to whole 4*4 matrix
     // lambda function for check data
     auto checkdata = [&](hamilt::AtomPair<double>& ap_in) {
@@ -671,4 +671,4 @@ int main(int argc, char** argv)
 #endif
 
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_complex.cpp b/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_complex.cpp
index dbc0088ecd..d55feb73c7 100644
--- a/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_complex.cpp
+++ b/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_complex.cpp
@@ -535,7 +535,7 @@ TEST_F(HContainerTest, atompair_funcs)
     // get_matrix_value will use global2local_row and global2local_col in Parallel_Orbitals
     // so we need to set them
     std::ofstream ofs("test_hcontainer_complex.log");
-    PO.set_global2local(4, 4, false, ofs);
+    PO.set_serial(4, 4);
     auto checkdata = [&](hamilt::AtomPair<std::complex<double>>& ap_in) {
         auto data_ij4 = ap_in.get_matrix_values();
         int* tmp_index = std::get<0>(data_ij4).data();
@@ -574,4 +574,4 @@ int main(int argc, char** argv)
 #endif
 
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_output.cpp b/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_output.cpp
index 5e37a17bcd..d9f0122e02 100644
--- a/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_output.cpp
+++ b/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_output.cpp
@@ -76,7 +76,7 @@ TEST_F(OutputHContainerTest, Write)
     ParaV.nrow = 4;
     ParaV.ncol = 4;
     std::ofstream ofs("output_hcontainer.log");
-    ParaV.set_global2local(4, 4, false, ofs);
+    ParaV.set_serial(4, 4);
     // std::cout << "ParaV.global2local_row = " << ParaV.global2local_row(0) << " " << ParaV.global2local_row(1) << " "
     //           << ParaV.global2local_row(2) << " " << ParaV.global2local_row(3) << std::endl;
     // std::cout << "ParaV.global2local_col = " << ParaV.global2local_col(0) << " " << ParaV.global2local_col(1) << " "
@@ -160,4 +160,4 @@ TEST_F(OutputHContainerTest, Write)
     EXPECT_THAT(output, testing::HasSubstr(" 5.00e+00 6.00e+00 1.00e+01"));
     EXPECT_THAT(output, testing::HasSubstr(" 2 3 3"));
     EXPECT_THAT(output, testing::HasSubstr(" 0 0 0 2 3"));
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_readCSR.cpp b/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_readCSR.cpp
index 8a4affd1f6..7d7ce6eb26 100644
--- a/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_readCSR.cpp
+++ b/source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_readCSR.cpp
@@ -73,7 +73,7 @@ TEST_F(ReadHContainerTest, ReadAndOutputHContainer)
     // construct paraV
     Parallel_Orbitals paraV;
     std::ofstream ofs("test.log");
-    paraV.set_global2local(nlocal, nlocal, false, ofs);
+    paraV.set_serial(nlocal, nlocal);
     ofs.close();
     remove("test.log");
     paraV.set_atomic_trace(ucell->get_iat2iwt(), ucell->nat, nlocal);
diff --git a/source/module_hamilt_lcao/module_hcontainer/test/test_transfer.cpp b/source/module_hamilt_lcao/module_hcontainer/test/test_transfer.cpp
index 40a7e387fe..78aa125613 100644
--- a/source/module_hamilt_lcao/module_hcontainer/test/test_transfer.cpp
+++ b/source/module_hamilt_lcao/module_hcontainer/test/test_transfer.cpp
@@ -56,18 +56,12 @@ class TransferTest : public ::testing::Test
 #ifdef __MPI
     void init_parav()
     {
+        int nb = 2;
         int global_row = test_size * test_nw;
         int global_col = test_size * test_nw;
         std::ofstream ofs_running;
         paraV = new Parallel_Orbitals();
-        paraV->set_block_size(2/* nb_2d set to be 2*/);
-        paraV->set_proc_dim(dsize, 0);
-        paraV->mpi_create_cart(MPI_COMM_WORLD);
-        paraV->set_local2global(global_row, global_col, ofs_running, ofs_running);
-        int lr = paraV->get_row_size();
-        int lc = paraV->get_col_size();
-        paraV->set_desc(global_row, global_col, lr);
-        paraV->set_global2local(global_row, global_col, true, ofs_running);
+        paraV->init(global_row, global_col, nb, MPI_COMM_WORLD);
         paraV->set_atomic_trace(ucell.get_iat2iwt(), test_size, global_row);
     }
 #else
@@ -237,7 +231,7 @@ TEST_F(TransferTest, paraToSerial)
     // now assume the HR_serial is empty and do the same test
     Parallel_Orbitals serialV;
     std::ofstream ofs("test.log");
-    serialV.set_global2local(test_size*test_nw, test_size*test_nw, false, ofs);
+    serialV.set_serial(test_size*test_nw, test_size*test_nw);
     serialV.set_atomic_trace(ucell.get_iat2iwt(), test_size, test_size*test_nw);
     hamilt::HContainer<double>* HR_serial2;
     if(my_rank == 0)
@@ -418,4 +412,4 @@ int main(int argc, char** argv)
     MPI_Finalize();
 #endif
     return result;
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/module_tddft/bandenergy.h b/source/module_hamilt_lcao/module_tddft/bandenergy.h
index e7f361f658..fb8c989d8b 100644
--- a/source/module_hamilt_lcao/module_tddft/bandenergy.h
+++ b/source/module_hamilt_lcao/module_tddft/bandenergy.h
@@ -8,6 +8,8 @@
 
 #include "module_basis/module_ao/parallel_orbitals.h"
 
+#include <complex>
+
 namespace module_tddft
 {
 #ifdef __MPI
@@ -29,4 +31,4 @@ void compute_ekb(const Parallel_Orbitals* pv,
                  double* ekb);
 #endif
 } // namespace module_tddft
-#endif
\ No newline at end of file
+#endif
diff --git a/source/module_hamilt_lcao/module_tddft/middle_hamilt.h b/source/module_hamilt_lcao/module_tddft/middle_hamilt.h
index 5e35f3bfa6..ac230b6444 100644
--- a/source/module_hamilt_lcao/module_tddft/middle_hamilt.h
+++ b/source/module_hamilt_lcao/module_tddft/middle_hamilt.h
@@ -7,6 +7,7 @@
 #define MIDDLE_HAMILT_H
 
 #include "module_basis/module_ao/parallel_orbitals.h"
+#include <complex>
 
 namespace module_tddft
 {
@@ -33,4 +34,4 @@ void half_Hmatrix(const Parallel_Orbitals* pv,
 #endif
 } // namespace module_tddft
 
-#endif
\ No newline at end of file
+#endif
diff --git a/source/module_hamilt_lcao/module_tddft/norm_psi.h b/source/module_hamilt_lcao/module_tddft/norm_psi.h
index 69287a1532..d3c41978bb 100644
--- a/source/module_hamilt_lcao/module_tddft/norm_psi.h
+++ b/source/module_hamilt_lcao/module_tddft/norm_psi.h
@@ -7,6 +7,7 @@
 #define NORM_PSI_H
 
 #include "module_basis/module_ao/parallel_orbitals.h"
+#include <complex>
 
 namespace module_tddft
 {
@@ -32,4 +33,4 @@ void norm_psi(const Parallel_Orbitals* pv,
 #endif
 } // namespace module_tddft
 
-#endif
\ No newline at end of file
+#endif
diff --git a/source/module_hamilt_lcao/module_tddft/propagator.h b/source/module_hamilt_lcao/module_tddft/propagator.h
index a156ca1355..4123d1447d 100644
--- a/source/module_hamilt_lcao/module_tddft/propagator.h
+++ b/source/module_hamilt_lcao/module_tddft/propagator.h
@@ -7,6 +7,7 @@
 #define PROPAGATOR_H
 
 #include "module_basis/module_ao/parallel_orbitals.h"
+#include <complex>
 
 namespace module_tddft
 {
@@ -97,4 +98,4 @@ class Propagator
 };
 } // namespace module_tddft
 
-#endif
\ No newline at end of file
+#endif
diff --git a/source/module_hamilt_lcao/module_tddft/test/bandenergy_test.cpp b/source/module_hamilt_lcao/module_tddft/test/bandenergy_test.cpp
index d6fe359389..2c164f50d0 100644
--- a/source/module_hamilt_lcao/module_tddft/test/bandenergy_test.cpp
+++ b/source/module_hamilt_lcao/module_tddft/test/bandenergy_test.cpp
@@ -20,8 +20,6 @@
 
 #define doublethreshold 1e-8
 double module_tddft::Evolve_elec::td_print_eij = -1;
-Parallel_2D::Parallel_2D(){}
-Parallel_2D::~Parallel_2D(){}
 Parallel_Orbitals::Parallel_Orbitals()
 {
 }
diff --git a/source/module_hamilt_lcao/module_tddft/test/middle_hamilt_test.cpp b/source/module_hamilt_lcao/module_tddft/test/middle_hamilt_test.cpp
index 66ec032148..5cda7e4aba 100644
--- a/source/module_hamilt_lcao/module_tddft/test/middle_hamilt_test.cpp
+++ b/source/module_hamilt_lcao/module_tddft/test/middle_hamilt_test.cpp
@@ -18,8 +18,6 @@
  */
 
 #define doublethreshold 1e-8
-Parallel_2D::Parallel_2D(){}
-Parallel_2D::~Parallel_2D(){}
 Parallel_Orbitals::Parallel_Orbitals()
 {
 }
diff --git a/source/module_hamilt_lcao/module_tddft/test/norm_psi_test.cpp b/source/module_hamilt_lcao/module_tddft/test/norm_psi_test.cpp
index 237426a30d..d92817af90 100644
--- a/source/module_hamilt_lcao/module_tddft/test/norm_psi_test.cpp
+++ b/source/module_hamilt_lcao/module_tddft/test/norm_psi_test.cpp
@@ -18,8 +18,6 @@
  */
 
 #define doublethreshold 1e-8
-Parallel_2D::Parallel_2D(){}
-Parallel_2D::~Parallel_2D(){}
 Parallel_Orbitals::Parallel_Orbitals()
 {
 }
diff --git a/source/module_hamilt_lcao/module_tddft/test/propagator_test1.cpp b/source/module_hamilt_lcao/module_tddft/test/propagator_test1.cpp
index 12abc8b770..415becbe70 100644
--- a/source/module_hamilt_lcao/module_tddft/test/propagator_test1.cpp
+++ b/source/module_hamilt_lcao/module_tddft/test/propagator_test1.cpp
@@ -19,8 +19,6 @@
 
 Input INPUT;
 #define doublethreshold 1e-8
-Parallel_2D::Parallel_2D(){}
-Parallel_2D::~Parallel_2D(){}
 Parallel_Orbitals::Parallel_Orbitals()
 {
 }
@@ -112,4 +110,4 @@ TEST(PropagatorTest, testPropagatorCN)
     delete[] U_operator;
     delete[] Htmp;
     delete[] Stmp;
-}
\ No newline at end of file
+}
diff --git a/source/module_hamilt_lcao/module_tddft/test/upsi_test1.cpp b/source/module_hamilt_lcao/module_tddft/test/upsi_test1.cpp
index bc308d32b6..2782db5376 100644
--- a/source/module_hamilt_lcao/module_tddft/test/upsi_test1.cpp
+++ b/source/module_hamilt_lcao/module_tddft/test/upsi_test1.cpp
@@ -17,8 +17,6 @@
  *   - upsi
  *     - apply U_operator to the wave function of the previous step for new wave function.
  */
-Parallel_2D::Parallel_2D(){}
-Parallel_2D::~Parallel_2D(){}
 Parallel_Orbitals::Parallel_Orbitals()
 {
 }
diff --git a/source/module_hamilt_lcao/module_tddft/upsi.h b/source/module_hamilt_lcao/module_tddft/upsi.h
index 7233eee77c..3d703de5f8 100644
--- a/source/module_hamilt_lcao/module_tddft/upsi.h
+++ b/source/module_hamilt_lcao/module_tddft/upsi.h
@@ -8,6 +8,7 @@
 #define UPSI_H
 
 #include "module_basis/module_ao/parallel_orbitals.h"
+#include <complex>
 
 namespace module_tddft
 {
@@ -34,4 +35,4 @@ void upsi(const Parallel_Orbitals* pv,
 #endif
 } // namespace module_tddft
 
-#endif
\ No newline at end of file
+#endif
diff --git a/source/module_io/single_R_io.h b/source/module_io/single_R_io.h
index 76abc590ad..b8661a437b 100644
--- a/source/module_io/single_R_io.h
+++ b/source/module_io/single_R_io.h
@@ -2,6 +2,7 @@
 #define SINGLE_R_IO_H
 
 #include "module_basis/module_ao/parallel_orbitals.h"
+#include <map>
 
 namespace ModuleIO
 {
diff --git a/source/module_io/test/single_R_io_test.cpp b/source/module_io/test/single_R_io_test.cpp
index 7bfc74f8b2..779e36626f 100644
--- a/source/module_io/test/single_R_io_test.cpp
+++ b/source/module_io/test/single_R_io_test.cpp
@@ -11,22 +11,17 @@
  *   - ModuleIO::output_single_R
  *     - output single R data
  */
-Parallel_2D::Parallel_2D(){}
-Parallel_2D::~Parallel_2D(){}
 Parallel_Orbitals::Parallel_Orbitals()
 {
-    global2local_row_ = nullptr;
 }
 
 Parallel_Orbitals::~Parallel_Orbitals()
 {
-    delete[] global2local_row_;
 }
 
-void Parallel_2D::set_global2local(const int& M_A, const int& N_A,
-    const bool& div_2d, std::ofstream& ofs_running)
+void Parallel_2D::set_serial(const int M_A, const int N_A)
 {
-    this->global2local_row_ = new int[M_A];
+    this->global2local_row_.resize(M_A);
     this->global2local_row_[0] = 0;
     this->global2local_row_[1] = 1;
     this->global2local_row_[2] = -1;
@@ -47,7 +42,7 @@ TEST(ModuleIOTest, OutputSingleR)
     const bool binary = false;
     Parallel_Orbitals pv;
     GlobalV::NLOCAL = 5;
-    pv.set_global2local(GlobalV::NLOCAL, GlobalV::NLOCAL, false, ofs);
+    pv.set_serial(GlobalV::NLOCAL, GlobalV::NLOCAL);
     std::map<size_t, std::map<size_t, double>> XR = {
         {0, {{1, 0.5}, {3, 0.3}}},
         {1, {{0, 0.2}, {2, 0.4}}},
diff --git a/source/module_io/test_serial/read_wfc_nao_test.cpp b/source/module_io/test_serial/read_wfc_nao_test.cpp
index 5d6dfe8743..8ff555f8fa 100644
--- a/source/module_io/test_serial/read_wfc_nao_test.cpp
+++ b/source/module_io/test_serial/read_wfc_nao_test.cpp
@@ -4,8 +4,6 @@
 #include "module_basis/module_ao/parallel_orbitals.h"
 
 //write mock function for Parallel_Orbitals
-Parallel_2D::Parallel_2D(){}
-Parallel_2D::~Parallel_2D(){}
 Parallel_Orbitals::Parallel_Orbitals() {}
 Parallel_Orbitals::~Parallel_Orbitals(){}
 //define a mock derived class of class ElecState
diff --git a/source/module_io/write_Vxc.hpp b/source/module_io/write_Vxc.hpp
index e384daf45e..99cc5a31d1 100644
--- a/source/module_io/write_Vxc.hpp
+++ b/source/module_io/write_Vxc.hpp
@@ -41,19 +41,9 @@ namespace ModuleIO
     {
         std::ofstream ofs;
 #ifdef __MPI
-        int dsize=0;
-        MPI_Comm_size(MPI_COMM_WORLD, &dsize);
-        p2d.set_block_size(pv.nb);
-        p2d.set_proc_dim(dsize);
-        p2d.comm_2D = pv.comm_2D;
-        p2d.blacs_ctxt = pv.blacs_ctxt;
-        p2d.set_local2global(nbands, nbands, ofs, ofs);
-        p2d.set_desc(nbands, nbands, p2d.get_row_size(), false);
-        p2d.set_global2local(nbands, nbands, true, ofs);
+        p2d.set(nbands, nbands, pv.nb, pv.comm_2D, pv.blacs_ctxt);
 #else
-        p2d.set_proc_dim(1);
         p2d.set_serial(nbands, nbands);
-        p2d.set_global2local(nbands, nbands, false, ofs);
 #endif
     }