refactor: remove the dependency of write_dmr on globalV::NLOCAL

source/module_esolver/esolver_ks_lcao.cpp

@@ -390,7 +390,7 @@ void ESolver_KS_LCAO<TK, TR>::cal_stress(ModuleBase::matrix& stress)
 //! mohan add 2024-05-11
 //------------------------------------------------------------------------------
 template <typename TK, typename TR>
-void ESolver_KS_LCAO<TK, TR>::after_all_runners(void)
+void ESolver_KS_LCAO<TK, TR>::after_all_runners()
 {
     ModuleBase::TITLE("ESolver_KS_LCAO", "after_all_runners");
     ModuleBase::timer::tick("ESolver_KS_LCAO", "after_all_runners");
@@ -1061,7 +1061,7 @@ void ESolver_KS_LCAO<TK, TR>::after_scf(const int istep)
 
     // 2) write density matrix for sparse matrix
     ModuleIO::write_dmr(dynamic_cast<const elecstate::ElecStateLCAO<TK>*>(this->pelec)->get_DM()->get_DMR_vector(),
-                        GlobalV::NLOCAL,
+                        this->orb_con.ParaV,
                         INPUT.out_dm1,
                         false,
                         GlobalV::out_app_flag,

source/module_hamilt_lcao/module_hcontainer/output_hcontainer.cpp

@@ -12,9 +12,11 @@ namespace hamilt
  * @attention ofs should be open outside of this interface
  */
 template <typename T>
-Output_HContainer<T>::Output_HContainer(hamilt::HContainer<T>* hcontainer, const int nrow, const int ncol,
-                                        std::ostream& ofs, double sparse_threshold, int precision)
-    : _hcontainer(hcontainer), _nrow(nrow), _ncol(ncol), _ofs(ofs), _sparse_threshold(sparse_threshold), _precision(precision)
+Output_HContainer<T>::Output_HContainer(hamilt::HContainer<T>* hcontainer,
+                                        std::ostream& ofs,
+                                        double sparse_threshold,
+                                        int precision)
+    : _hcontainer(hcontainer), _ofs(ofs), _sparse_threshold(sparse_threshold), _precision(precision)
 {
 }
 
@@ -56,7 +58,8 @@ template <typename T>
 void Output_HContainer<T>::write_single_R(int rx, int ry, int rz)
 {
     this->_hcontainer->fix_R(rx, ry, rz);
-    ModuleIO::SparseMatrix<T> sparse_matrix = ModuleIO::SparseMatrix<T>(_nrow, _ncol);
+    ModuleIO::SparseMatrix<T> sparse_matrix
+        = ModuleIO::SparseMatrix<T>(_hcontainer->get_nbasis(), _hcontainer->get_nbasis());
     sparse_matrix.setSparseThreshold(this->_sparse_threshold);
     for (int iap = 0; iap < this->_hcontainer->size_atom_pairs(); ++iap)
     {

source/module_hamilt_lcao/module_hcontainer/output_hcontainer.h

@@ -13,8 +13,7 @@ template <typename T>
 class Output_HContainer
 {
   public:
-    Output_HContainer(hamilt::HContainer<T>* hcontainer, const int nrow, const int ncol, std::ostream& ofs,
-                      double sparse_threshold, int precision);
+    Output_HContainer(hamilt::HContainer<T>* hcontainer, std::ostream& ofs, double sparse_threshold, int precision);
     // write the matrices of all R vectors to the output stream
     void write();
 
@@ -32,8 +31,6 @@ class Output_HContainer
 
   private:
     hamilt::HContainer<T>* _hcontainer;
-    int _nrow=0; 
-    int _ncol=0;
     std::ostream& _ofs;
     double _sparse_threshold;
     int _precision;

source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_output.cpp

@@ -131,7 +131,7 @@ TEST_F(OutputHContainerTest, Write)
         HR.unfix_R();
     }
     double sparse_threshold = 0.1;
-    hamilt::Output_HContainer<double> output_HR(&HR, ParaV.nrow, ParaV.ncol, std::cout, sparse_threshold, 2);
+    hamilt::Output_HContainer<double> output_HR(&HR, std::cout, sparse_threshold, 2);
     // the first R
     testing::internal::CaptureStdout();
     output_HR.write(0, 1, 1);

source/module_hamilt_lcao/module_hcontainer/test/test_hcontainer_readCSR.cpp

@@ -124,7 +124,7 @@ TEST_F(ReadHContainerTest, ReadAndOutputHContainer)
     // output SR
     std::ofstream ofs_out("SR.out");
     double sparse_threshold = 1e-10;
-    hamilt::Output_HContainer<double> output_SR(&SR, paraV.nrow, paraV.ncol, ofs_out, sparse_threshold, 8);
+    hamilt::Output_HContainer<double> output_SR(&SR, ofs_out, sparse_threshold, 8);
     // std::cout << SR.size_R_loop() << std::endl;
     // output_SR.write(-2, -1, 0);
     ofs_out << "STEP: " << 0 << std::endl;

source/module_io/write_dmr.cpp

@@ -33,24 +33,28 @@ std::string dmr_gen_fname(const int out_type, const int ispin, const bool append
     return fname;
 }
 
-void write_dmr_csr(std::string& fname, hamilt::HContainer<double>* dm_serial, const int nbasis, const int istep)
+void write_dmr_csr(std::string& fname, hamilt::HContainer<double>* dm_serial, const int istep)
 {
     // write the head: ION step number, basis number and R loop number
     std::ofstream ofs(fname, std::ios::app);
     ofs << "STEP: " << istep << std::endl;
-    ofs << "Matrix Dimension of DM(R): " << nbasis << std::endl;
+    ofs << "Matrix Dimension of DM(R): " << dm_serial->get_nbasis() << std::endl;
     ofs << "Matrix number of DM(R): " << dm_serial->size_R_loop() << std::endl;
 
     // write HR_serial to ofs
     double sparse_threshold = 1e-10;
     int precision = 8;
-    hamilt::Output_HContainer<double> out_dm(dm_serial, nbasis, nbasis, ofs, sparse_threshold, precision);
+    hamilt::Output_HContainer<double> out_dm(dm_serial, ofs, sparse_threshold, precision);
     out_dm.write();
     ofs.close();
 }
 
-void write_dmr(const std::vector<hamilt::HContainer<double>*> dmr, const int nbasis,const bool out_csr, const bool out_npz,
-               const bool append, const int istep)
+void write_dmr(const std::vector<hamilt::HContainer<double>*> dmr,
+               const Parallel_2D& paraV,
+               const bool out_csr,
+               const bool out_npz,
+               const bool append,
+               const int istep)
 {
     if (!out_csr && !out_npz)
     {
@@ -60,30 +64,29 @@ void write_dmr(const std::vector<hamilt::HContainer<double>*> dmr, const int nba
 
     for (int ispin = 0; ispin < dmr.size(); ispin++)
     {
-        // gather the parallel matrix to serial matrix
+        if (out_csr)
+        {
+            int nbasis = dmr[ispin]->get_nbasis();
+            // gather the parallel matrix to serial matrix
 #ifdef __MPI
-        Parallel_Orbitals serialV;
-        serialV.set_serial(nbasis, nbasis);
-        serialV.set_atomic_trace(GlobalC::ucell.get_iat2iwt(), GlobalC::ucell.nat, nbasis);
-        hamilt::HContainer<double> dm_serial(&serialV);
-        hamilt::gatherParallels(*dmr[ispin], &dm_serial, 0);
+            Parallel_Orbitals serialV;
+            serialV.init(nbasis, nbasis, nbasis, paraV.comm_2D);
+            serialV.set_serial(nbasis, nbasis);
+            serialV.set_atomic_trace(GlobalC::ucell.get_iat2iwt(), GlobalC::ucell.nat, nbasis);
+            hamilt::HContainer<double> dm_serial(&serialV);
+            hamilt::gatherParallels(*dmr[ispin], &dm_serial, 0);
 #else
-        hamilt::HContainer<double> dm_serial(*dmr[ispin]);
+            hamilt::HContainer<double> dm_serial(*dmr[ispin]);
 #endif
-
-        if(out_csr)
-        {
             if (GlobalV::MY_RANK == 0)
             {
                 std::string fname = GlobalV::global_out_dir + dmr_gen_fname(1, ispin, append, istep);
-                write_dmr_csr(fname, &dm_serial, nbasis, istep);
+                write_dmr_csr(fname, &dm_serial, istep);
             }
         }
 
-        if(out_npz)
+        if (out_npz)
         {
-            // std::string fname = GlobalV::global_out_dir + dmr_gen_fname(2, ispin, append, istep);
-            // write_dmr_npz(fname, *dm_serial, istep);
         }
     }
 }

source/module_io/write_dmr.h

@@ -17,34 +17,33 @@ namespace ModuleIO
  * @param istep    The ION step (default: -1), starting from 0.
  * @return         The generated filename as a string.
  */
-std::string dmr_gen_fname(const int out_type, const int ispin, const bool append = true,
-                          const int istep = -1);
-
+std::string dmr_gen_fname(const int out_type, const int ispin, const bool append = true, const int istep = -1);
 
 /**
  * Writes HContainer to a csr file.
  *
  * @param fname The name of the file to write the CSR representation to.
  * @param dm_serial A pointer to the Hamiltonian container.
- * @param nbasis The number of basis functions.
  * @param istep The current step number.
  */
-void write_dmr_csr(std::string& fname, hamilt::HContainer<double>* dm_serial, const int nbasis, const int istep);
+void write_dmr_csr(std::string& fname, hamilt::HContainer<double>* dm_serial, const int istep);
 
 /**
  * Writes DMR to a file.
  *
- * @param dm The 2D block parallel matrix representing the density matrix.
- * @param nbasis The number of basis functions.
+ * @param dmr The 2D block parallel matrix representing the density matrix. The first dimension is the spin index.
+ * @param paraV The parallel 2D object.
  * @param out_type The output file type. 1: csr, 2: npz.
  * @param sparse Whether output the sparse DM.
  * @param ispin The spin index, starting from 0.
  * @param append Whether to append the data to an existing file or create a new file. The file name is related to this
  * flag.
  * @param istep The ION step, starting from 0.
- * @param pv The Parallel_Orbitals object.
  */
-void write_dmr(const std::vector<hamilt::HContainer<double>*> dmr, const int nbasis, const bool out_csr, const bool out_npz,
+void write_dmr(const std::vector<hamilt::HContainer<double>*> dmr,
+               const Parallel_2D& paraV,
+               const bool out_csr,
+               const bool out_npz,
                const bool append,
                const int istep);
 } // namespace ModuleIO