Refactor(io): modularize sparse matrix writer interface

source/source_io/module_hs/cal_pLpR.cpp

@@ -180,8 +180,13 @@ ModuleIO::AngularMomentumCalculator::AngularMomentumCalculator(
     const int rank)
 {
 
-    // ofs_running
     this->ofs_ = ptr_log;
+    if (this->ofs_ == nullptr)
+    {
+        this->fallback_ofs_.open("/dev/null");
+        this->ofs_ = &this->fallback_ofs_;
+    }
+
     *ofs_ << "\n\n\n\n";
     *ofs_ << " >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>" << std::endl;
     *ofs_ << " |                                                                    |" << std::endl;
@@ -259,7 +264,7 @@ void ModuleIO::AngularMomentumCalculator::kernel(
     const char dir,
     const int precision)
 {
-    if (!ofs->is_open())
+    if (ofs == nullptr || !ofs->is_open())
     {
         return;
     }
@@ -381,4 +386,4 @@ void ModuleIO::AngularMomentumCalculator::calculate(
         this->kernel(&ofout, ucell, d, precision);
         ofout.close();
     }
-}
+}

source/source_io/module_hs/cal_pLpR.h

@@ -88,6 +88,7 @@
 #include <map>
 #include <tuple>
 #include <complex>
+#include <fstream>
 #include <memory>
 #include "source_cell/unitcell.h"
 #include "source_basis/module_nao/two_center_integrator.h"
@@ -218,8 +219,9 @@ namespace ModuleIO
                            const int rank = 0);
 
         private:
+            std::ofstream fallback_ofs_;
             // ofsrunning
-            std::ofstream* ofs_;
+            std::ofstream* ofs_ = nullptr;
             // the two-center-integrator
             std::unique_ptr<TwoCenterIntegrator> calculator_;
             // the spherical bessel transformer
@@ -246,4 +248,4 @@ namespace ModuleIO
                         const char dir = 'x',
                         const int precision = 10);
     };
-} // namespace ModuleIO
+} // namespace ModuleIO

source/source_io/module_hs/cal_r_overlap_R.cpp

@@ -628,6 +628,11 @@ void cal_r_overlap_R::out_rR(const UnitCell& ucell, const Grid_Driver& gd, const
     ModuleBase::Vector3<double> origin_point(0.0, 0.0, 0.0);
     double factor = sqrt(ModuleBase::FOUR_PI / 3.0);
     int output_R_number = 0;
+    ModuleIO::SparseWriteOptions single_R_options;
+    single_R_options.threshold = sparse_threshold;
+    single_R_options.binary = binary;
+    single_R_options.reduce = true;
+    single_R_options.temp_dir = PARAM.globalv.global_out_dir;
 
     std::stringstream tem1;
     tem1 << PARAM.globalv.global_out_dir << "tmp-rr.csr";
@@ -789,9 +794,8 @@ void cal_r_overlap_R::out_rR(const UnitCell& ucell, const Grid_Driver& gd, const
                 {
                     ModuleIO::output_single_R(ofs_tem1,
                                               psi_r_psi_sparse[direction],
-                                              sparse_threshold,
-                                              binary,
-                                              *(this->ParaV));
+                                              *(this->ParaV),
+                                              single_R_options);
                 }
                 else
                 {
@@ -875,6 +879,11 @@ void cal_r_overlap_R::out_rR_other(const UnitCell& ucell, const int& istep, cons
     double factor = sqrt(ModuleBase::FOUR_PI / 3.0);
     int output_R_number = output_R_coor.size();
     int step = istep;
+    ModuleIO::SparseWriteOptions single_R_options;
+    single_R_options.threshold = sparse_threshold;
+    single_R_options.binary = binary;
+    single_R_options.reduce = true;
+    single_R_options.temp_dir = PARAM.globalv.global_out_dir;
 
     std::ofstream out_r;
     std::stringstream ssr;
@@ -1059,7 +1068,10 @@ void cal_r_overlap_R::out_rR_other(const UnitCell& ucell, const int& istep, cons
 
             if (rR_nonzero_num[direction])
             {
-                ModuleIO::output_single_R(out_r, psi_r_psi_sparse[direction], sparse_threshold, binary, *(this->ParaV));
+                ModuleIO::output_single_R(out_r,
+                                          psi_r_psi_sparse[direction],
+                                          *(this->ParaV),
+                                          single_R_options);
             }
             else
             {

source/source_io/module_hs/single_R_io.cpp

@@ -1,9 +1,14 @@
 #include "single_R_io.h"
 #include "source_base/parallel_reduce.h"
-#include "source_io/module_parameter/parameter.h"
 #include "source_base/global_function.h"
 #include "source_base/global_variable.h"
 
+#include <complex>
+#include <cstdio>
+#include <iomanip>
+#include <sstream>
+#include <vector>
+
 inline void write_data(std::ofstream& ofs, const double& data)
 {
     ofs << " " << std::fixed << std::scientific << std::setprecision(16) << data;
@@ -15,25 +20,30 @@ inline void write_data(std::ofstream& ofs, const std::complex<double>& data)
 
 template<typename T>
 void ModuleIO::output_single_R(std::ofstream& ofs,
-    const std::map<size_t, std::map<size_t, T>>& XR,
-    const double& sparse_threshold,
-    const bool& binary,
+    const SparseRBlock<T>& XR,
     const Parallel_Orbitals& pv,
-    const bool& reduce)
+    const SparseWriteOptions& options)
 {
-    T* line = nullptr;
+    const int nlocal = pv.get_global_row_size();
+    if (nlocal <= 0)
+    {
+        ModuleBase::WARNING_QUIT("ModuleIO::output_single_R",
+                                 "Parallel_Orbitals global row size must be positive.");
+    }
+
     std::vector<long long> indptr;
-    indptr.reserve(PARAM.globalv.nlocal + 1);
+    indptr.reserve(nlocal + 1);
     indptr.push_back(0);
 
     std::stringstream tem1;
-    tem1 << PARAM.globalv.global_out_dir << std::to_string(GlobalV::DRANK) + "temp_sparse_indices.dat";
+    tem1 << options.temp_dir << std::to_string(GlobalV::DRANK)
+         << "temp_sparse_indices.dat";
     std::ofstream ofs_tem1;
     std::ifstream ifs_tem1;
 
-    if (!reduce || GlobalV::DRANK == 0)
+    if (!options.reduce || GlobalV::DRANK == 0)
     {
-        if (binary)
+        if (options.binary)
         {
             ofs_tem1.open(tem1.str().c_str(), std::ios::binary);
         }
@@ -43,12 +53,12 @@ void ModuleIO::output_single_R(std::ofstream& ofs,
         }
     }
 
-    line = new T[PARAM.globalv.nlocal];
-    for(int row = 0; row < PARAM.globalv.nlocal; ++row)
+    std::vector<T> line(nlocal);
+    for(int row = 0; row < nlocal; ++row)
     {
-        ModuleBase::GlobalFunc::ZEROS(line, PARAM.globalv.nlocal);
+        ModuleBase::GlobalFunc::ZEROS(line.data(), nlocal);
 
-        if (!reduce || pv.global2local_row(row) >= 0)
+        if (!options.reduce || pv.global2local_row(row) >= 0)
         {
             auto iter = XR.find(row);
             if (iter != XR.end())
@@ -60,19 +70,19 @@ void ModuleIO::output_single_R(std::ofstream& ofs,
             }
         }
 
-		if (reduce) 
-		{
-			Parallel_Reduce::reduce_all(line, PARAM.globalv.nlocal);
-		}
+        if (options.reduce)
+        {
+            Parallel_Reduce::reduce_all(line.data(), nlocal);
+        }
 
-        if (!reduce || GlobalV::DRANK == 0)
+        if (!options.reduce || GlobalV::DRANK == 0)
         {
             long long nonzeros_count = 0;
-            for (int col = 0; col < PARAM.globalv.nlocal; ++col)
+            for (int col = 0; col < nlocal; ++col)
             {
-                if (std::abs(line[col]) > sparse_threshold)
+                if (std::abs(line[col]) > options.threshold)
                 {
-                    if (binary)
+                    if (options.binary)
                     {
                         ofs.write(reinterpret_cast<char*>(&line[col]), sizeof(T));
                         ofs_tem1.write(reinterpret_cast<char *>(&col), sizeof(int));
@@ -93,11 +103,9 @@ void ModuleIO::output_single_R(std::ofstream& ofs,
         }
     }
 
-    delete[] line;
-
-    if (!reduce || GlobalV::DRANK == 0)
+    if (!options.reduce || GlobalV::DRANK == 0)
     {
-        if (binary)
+        if (options.binary)
         {
             ofs_tem1.close();
             ifs_tem1.open(tem1.str().c_str(), std::ios::binary);
@@ -128,15 +136,11 @@ void ModuleIO::output_single_R(std::ofstream& ofs,
 }
 
 template void ModuleIO::output_single_R<double>(std::ofstream& ofs,
-    const std::map<size_t, std::map<size_t, double>>& XR,
-    const double& sparse_threshold,
-    const bool& binary,
+    const SparseRBlock<double>& XR,
     const Parallel_Orbitals& pv,
-    const bool& reduce);
+    const SparseWriteOptions& options);
 
 template void ModuleIO::output_single_R<std::complex<double>>(std::ofstream& ofs,
-    const std::map<size_t, std::map<size_t, std::complex<double>>>& XR,
-    const double& sparse_threshold,
-    const bool& binary,
+    const SparseRBlock<std::complex<double>>& XR,
     const Parallel_Orbitals& pv,
-    const bool& reduce);
+    const SparseWriteOptions& options);

source/source_io/module_hs/single_R_io.h

@@ -1,18 +1,17 @@
 #ifndef SINGLE_R_IO_H
 #define SINGLE_R_IO_H
 
-#include "source_basis/module_ao/parallel_orbitals.h"
-#include <map>
+#include "write_HS_sparse.h"
+
+#include <fstream>
 
 namespace ModuleIO
 {
     template <typename T>
     void output_single_R(std::ofstream& ofs,
-        const std::map<size_t, std::map<size_t, T>>& XR,
-        const double& sparse_threshold,
-        const bool& binary,
+        const SparseRBlock<T>& XR,
         const Parallel_Orbitals& pv,
-        const bool& reduce = true);
+        const SparseWriteOptions& options);
 }
 
 #endif

source/source_io/module_hs/write_HS.hpp

@@ -280,11 +280,11 @@ void ModuleIO::save_mat(const int istep,
 #else
         if (app)
         {
-            std::ofstream out_matrix(filename.c_str(), std::ofstream::app);
+            out_matrix.open(filename.c_str(), std::ofstream::app);
         }
         else
         {
-            std::ofstream out_matrix(filename.c_str());
+            out_matrix.open(filename.c_str());
         }
 
         out_matrix << dim;

source/source_io/module_hs/write_HS_R.cpp

@@ -8,34 +8,6 @@
 #include "source_lcao/spar_st.h"
 #include "write_HS_sparse.h"
 
-namespace {
-// Helper: Convert sparse map to HContainer
-template <typename T>
-hamilt::HContainer<T>* sparse_map_to_hcontainer(
-    const std::map<Abfs::Vector3_Order<int>, std::map<size_t, std::map<size_t, T>>>& sparse_map,
-    const Parallel_Orbitals& pv,
-    const int nbasis)
-{
-    hamilt::HContainer<T>* hc = new hamilt::HContainer<T>(&pv);
-    hc->set_zero();
-
-    for (const auto& r_entry : sparse_map)
-    {
-        const auto& R = r_entry.first;
-        for (const auto& row_entry : r_entry.second)
-        {
-            const size_t row = row_entry.first;
-            for (const auto& col_entry : row_entry.second)
-            {
-                hc->set_value(R.x, R.y, R.z, row, col_entry.first, col_entry.second);
-            }
-        }
-    }
-
-    return hc;
-}
-} // anonymous namespace
-
 // if 'binary=true', output binary file.
 // The 'sparse_thr' is the accuracy of the sparse matrix.
 // If the absolute value of the matrix element is less than or equal to the
@@ -143,28 +115,28 @@ void ModuleIO::output_SR(Parallel_Orbitals& pv,
                           p_ham);
 
     const int istep = 0;
+    ModuleIO::SparseWriteOptions options;
+    options.filename = SR_filename;
+    options.label = "S";
+    options.threshold = sparse_thr;
+    options.binary = binary;
+    options.istep = istep;
+    options.reduce = true;
+    options.temp_dir = PARAM.globalv.global_out_dir;
 
     if (PARAM.inp.nspin == 4)
     {
         ModuleIO::save_sparse(HS_Arrays.SR_soc_sparse,
                               HS_Arrays.all_R_coor,
-                              sparse_thr,
-                              binary,
-                              SR_filename,
                               pv,
-                              "S",
-                              istep);
+                              options);
     }
     else
     {
         ModuleIO::save_sparse(HS_Arrays.SR_sparse,
                               HS_Arrays.all_R_coor,
-                              sparse_thr,
-                              binary,
-                              SR_filename,
                               pv,
-                              "S",
-                              istep);
+                              options);
     }
 
     sparse_format::destroy_HS_R_sparse(HS_Arrays);
@@ -206,15 +178,19 @@ void ModuleIO::output_TR(const int istep,
     }
 
     sparse_format::cal_TR(ucell, pv, HS_Arrays, grid, two_center_bundle, orb, sparse_thr);
+    ModuleIO::SparseWriteOptions options;
+    options.filename = sst.str();
+    options.label = "T";
+    options.threshold = sparse_thr;
+    options.binary = binary;
+    options.istep = istep;
+    options.reduce = true;
+    options.temp_dir = PARAM.globalv.global_out_dir;
 
     ModuleIO::save_sparse(HS_Arrays.TR_sparse,
                           HS_Arrays.all_R_coor,
-                          sparse_thr,
-                          binary,
-                          sst.str().c_str(),
                           pv,
-                          "T",
-                          istep);
+                          options);
 
     sparse_format::destroy_T_R_sparse(HS_Arrays);