diff --git a/source/source_io/module_hs/cal_pLpR.cpp b/source/source_io/module_hs/cal_pLpR.cpp index a9eb2223762..f98dc1dc474 100644 --- a/source/source_io/module_hs/cal_pLpR.cpp +++ b/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(); } -} \ No newline at end of file +} diff --git a/source/source_io/module_hs/cal_pLpR.h b/source/source_io/module_hs/cal_pLpR.h index 9ecd15a6120..c4694596cec 100644 --- a/source/source_io/module_hs/cal_pLpR.h +++ b/source/source_io/module_hs/cal_pLpR.h @@ -88,6 +88,7 @@ #include #include #include +#include #include #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 calculator_; // the spherical bessel transformer @@ -246,4 +248,4 @@ namespace ModuleIO const char dir = 'x', const int precision = 10); }; -} // namespace ModuleIO \ No newline at end of file +} // namespace ModuleIO diff --git a/source/source_io/module_hs/cal_r_overlap_R.cpp b/source/source_io/module_hs/cal_r_overlap_R.cpp index f75570a4116..7a26e2d7fc8 100644 --- a/source/source_io/module_hs/cal_r_overlap_R.cpp +++ b/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 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 { diff --git a/source/source_io/module_hs/single_R_io.cpp b/source/source_io/module_hs/single_R_io.cpp index bba94c345ee..f5e52a0c213 100644 --- a/source/source_io/module_hs/single_R_io.cpp +++ b/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 +#include +#include +#include +#include + 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& data) template void ModuleIO::output_single_R(std::ofstream& ofs, - const std::map>& XR, - const double& sparse_threshold, - const bool& binary, + const SparseRBlock& 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 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 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(&line[col]), sizeof(T)); ofs_tem1.write(reinterpret_cast(&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(std::ofstream& ofs, - const std::map>& XR, - const double& sparse_threshold, - const bool& binary, + const SparseRBlock& XR, const Parallel_Orbitals& pv, - const bool& reduce); + const SparseWriteOptions& options); template void ModuleIO::output_single_R>(std::ofstream& ofs, - const std::map>>& XR, - const double& sparse_threshold, - const bool& binary, + const SparseRBlock>& XR, const Parallel_Orbitals& pv, - const bool& reduce); + const SparseWriteOptions& options); diff --git a/source/source_io/module_hs/single_R_io.h b/source/source_io/module_hs/single_R_io.h index df6e0900a19..405a5bee4e2 100644 --- a/source/source_io/module_hs/single_R_io.h +++ b/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 +#include "write_HS_sparse.h" + +#include namespace ModuleIO { template void output_single_R(std::ofstream& ofs, - const std::map>& XR, - const double& sparse_threshold, - const bool& binary, + const SparseRBlock& XR, const Parallel_Orbitals& pv, - const bool& reduce = true); + const SparseWriteOptions& options); } #endif diff --git a/source/source_io/module_hs/write_HS.hpp b/source/source_io/module_hs/write_HS.hpp index 0d4950ef9c9..58e18a07554 100644 --- a/source/source_io/module_hs/write_HS.hpp +++ b/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; diff --git a/source/source_io/module_hs/write_HS_R.cpp b/source/source_io/module_hs/write_HS_R.cpp index 2a16db86d1f..be69653eac9 100644 --- a/source/source_io/module_hs/write_HS_R.cpp +++ b/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 -hamilt::HContainer* sparse_map_to_hcontainer( - const std::map, std::map>>& sparse_map, - const Parallel_Orbitals& pv, - const int nbasis) -{ - hamilt::HContainer* hc = new hamilt::HContainer(&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); diff --git a/source/source_io/module_hs/write_HS_sparse.cpp b/source/source_io/module_hs/write_HS_sparse.cpp index adcdd74951d..ed4a7387412 100644 --- a/source/source_io/module_hs/write_HS_sparse.cpp +++ b/source/source_io/module_hs/write_HS_sparse.cpp @@ -1,11 +1,117 @@ #include "write_HS_sparse.h" -#include "source_io/module_parameter/parameter.h" +#include "source_base/global_function.h" #include "source_base/parallel_reduce.h" #include "source_base/timer.h" +#include "source_io/module_parameter/parameter.h" #include "source_lcao/module_rt/td_info.h" #include "single_R_io.h" +#include +#include + +namespace +{ +template +std::vector count_nonzeros_by_R( + const ModuleIO::SparseRMatrix& smat, + const std::set& all_R_coor, + const bool reduce) +{ + std::vector nonzero_num(all_R_coor.size(), 0); + int count = 0; + for (const auto& R_coor: all_R_coor) + { + const auto iter = smat.find(R_coor); + if (iter != smat.end()) + { + for (const auto& row_loop: iter->second) + { + nonzero_num[count] += row_loop.second.size(); + } + } + ++count; + } + + if (reduce) + { + Parallel_Reduce::reduce_all(nonzero_num.data(), static_cast(nonzero_num.size())); + } + return nonzero_num; +} + +int count_output_R(const std::vector& nonzero_num) +{ + int output_R_number = 0; + for (const long long count: nonzero_num) + { + if (count != 0) + { + ++output_R_number; + } + } + return output_R_number; +} + +void open_sparse_file(std::ofstream& ofs, const ModuleIO::SparseWriteOptions& options) +{ + std::ios_base::openmode mode = std::ios::out; + if (options.binary) + { + mode |= std::ios::binary; + } + if (PARAM.inp.calculation == "md" && PARAM.inp.out_app_flag && options.istep) + { + mode |= std::ios::app; + } + ofs.open(options.filename.c_str(), mode); +} + +void write_sparse_header(std::ofstream& ofs, + const ModuleIO::SparseWriteOptions& options, + const int nlocal, + const int output_R_number) +{ + const int step = std::max(options.istep, 0); + if (options.binary) + { + ofs.write(reinterpret_cast(&step), sizeof(int)); + ofs.write(reinterpret_cast(&nlocal), sizeof(int)); + ofs.write(reinterpret_cast(&output_R_number), sizeof(int)); + } + else + { + ofs << "STEP: " << step << std::endl; + ofs << "Matrix Dimension of " + options.label + "(R): " << nlocal + << std::endl; + ofs << "Matrix number of " + options.label + "(R): " + << output_R_number << std::endl; + } +} + +void write_R_record(std::ofstream& ofs, + const ModuleIO::RCoordinate& R_coor, + const long long nonzero_count, + const bool binary) +{ + int dRx = R_coor.x; + int dRy = R_coor.y; + int dRz = R_coor.z; + if (binary) + { + const int count = static_cast(nonzero_count); + ofs.write(reinterpret_cast(&dRx), sizeof(int)); + ofs.write(reinterpret_cast(&dRy), sizeof(int)); + ofs.write(reinterpret_cast(&dRz), sizeof(int)); + ofs.write(reinterpret_cast(&count), sizeof(int)); + } + else + { + ofs << dRx << " " << dRy << " " << dRz << " " << nonzero_count + << std::endl; + } +} +} // namespace void ModuleIO::save_dH_sparse(const int& istep, const Parallel_Orbitals& pv, @@ -15,6 +121,11 @@ void ModuleIO::save_dH_sparse(const int& istep, const std::string& fileflag) { ModuleBase::TITLE("ModuleIO", "save_dH_sparse"); ModuleBase::timer::start("ModuleIO", "save_dH_sparse"); + SparseWriteOptions single_R_options; + single_R_options.threshold = sparse_thr; + single_R_options.binary = binary; + single_R_options.reduce = true; + single_R_options.temp_dir = PARAM.globalv.global_out_dir; auto& all_R_coor_ptr = HS_Arrays.all_R_coor; auto& output_R_coor_ptr = HS_Arrays.output_R_coor; @@ -307,45 +418,39 @@ void ModuleIO::save_dH_sparse(const int& istep, if (PARAM.inp.nspin != 4) { output_single_R(g1x[ispin], dHRx_sparse_ptr[ispin][R_coor], - sparse_thr, - binary, - pv); + pv, + single_R_options); } else { output_single_R(g1x[ispin], dHRx_soc_sparse_ptr[R_coor], - sparse_thr, - binary, - pv); + pv, + single_R_options); } } if (dHy_nonzero_num[ispin][count] > 0) { if (PARAM.inp.nspin != 4) { output_single_R(g1y[ispin], dHRy_sparse_ptr[ispin][R_coor], - sparse_thr, - binary, - pv); + pv, + single_R_options); } else { output_single_R(g1y[ispin], dHRy_soc_sparse_ptr[R_coor], - sparse_thr, - binary, - pv); + pv, + single_R_options); } } if (dHz_nonzero_num[ispin][count] > 0) { if (PARAM.inp.nspin != 4) { output_single_R(g1z[ispin], dHRz_sparse_ptr[ispin][R_coor], - sparse_thr, - binary, - pv); + pv, + single_R_options); } else { output_single_R(g1z[ispin], dHRz_soc_sparse_ptr[R_coor], - sparse_thr, - binary, - pv); + pv, + single_R_options); } } } @@ -380,108 +485,56 @@ void ModuleIO::save_dH_sparse(const int& istep, template void ModuleIO::save_sparse( - const std::map, - std::map>>& smat, - const std::set>& all_R_coor, - const double& sparse_thr, - const bool& binary, - const std::string& filename, + const SparseRMatrix& smat, + const std::set& all_R_coor, const Parallel_Orbitals& pv, - const std::string& label, - const int& istep, - const bool& reduce) { + const SparseWriteOptions& options) { ModuleBase::TITLE("ModuleIO", "save_sparse"); ModuleBase::timer::start("ModuleIO", "save_sparse"); - - int total_R_num = all_R_coor.size(); - std::vector nonzero_num(total_R_num, 0); - int count = 0; - for (auto& R_coor: all_R_coor) { - auto iter = smat.find(R_coor); - if (iter != smat.end()) { - for (auto& row_loop: iter->second) { - nonzero_num[count] += row_loop.second.size(); - } - } - ++count; - } - if (reduce) { - Parallel_Reduce::reduce_all(nonzero_num.data(), total_R_num); - } - - int output_R_number = 0; - for (int index = 0; index < total_R_num; ++index) { - if (nonzero_num[index] != 0) { - ++output_R_number; - } + const int nlocal = pv.get_global_row_size(); + if (nlocal <= 0) + { + ModuleBase::WARNING_QUIT("ModuleIO::save_sparse", + "Parallel_Orbitals global row size must be positive."); } - std::stringstream sss; - sss << filename; + const std::vector nonzero_num + = count_nonzeros_by_R(smat, all_R_coor, options.reduce); + const int output_R_number = count_output_R(nonzero_num); std::ofstream ofs; - if (!reduce || GlobalV::DRANK == 0) { - if (binary) { - int nlocal = PARAM.globalv.nlocal; - if (PARAM.inp.calculation == "md" && PARAM.inp.out_app_flag - && istep) { - ofs.open(sss.str().c_str(), std::ios::binary | std::ios::app); - } else { - ofs.open(sss.str().c_str(), std::ios::binary); - } - ofs.write(reinterpret_cast(0), sizeof(int)); - ofs.write(reinterpret_cast(&nlocal), sizeof(int)); - ofs.write(reinterpret_cast(&output_R_number), sizeof(int)); - } else { - if (PARAM.inp.calculation == "md" && PARAM.inp.out_app_flag - && istep) { - ofs.open(sss.str().c_str(), std::ios::app); - } else { - ofs.open(sss.str().c_str()); - } - ofs << "STEP: " << std::max(istep, 0) << std::endl; - ofs << "Matrix Dimension of " + label + "(R): " << PARAM.globalv.nlocal - << std::endl; - ofs << "Matrix number of " + label + "(R): " << output_R_number - << std::endl; - } + if (!options.reduce || GlobalV::DRANK == 0) + { + open_sparse_file(ofs, options); + write_sparse_header(ofs, options, nlocal, output_R_number); } - count = 0; - for (auto& R_coor: all_R_coor) { - int dRx = R_coor.x; - int dRy = R_coor.y; - int dRz = R_coor.z; - - if (nonzero_num[count] == 0) { + int count = 0; + for (const auto& R_coor: all_R_coor) + { + if (nonzero_num[count] == 0) + { count++; continue; } - if (!reduce || GlobalV::DRANK == 0) { - if (binary) { - ofs.write(reinterpret_cast(&dRx), sizeof(int)); - ofs.write(reinterpret_cast(&dRy), sizeof(int)); - ofs.write(reinterpret_cast(&dRz), sizeof(int)); - ofs.write(reinterpret_cast(&nonzero_num[count]), - sizeof(int)); - } else { - ofs << dRx << " " << dRy << " " << dRz << " " - << nonzero_num[count] << std::endl; - } + if (!options.reduce || GlobalV::DRANK == 0) + { + write_R_record(ofs, R_coor, nonzero_num[count], options.binary); } if (smat.count(R_coor)) { - output_single_R(ofs, smat.at(R_coor), sparse_thr, binary, pv, reduce); + output_single_R(ofs, smat.at(R_coor), pv, options); } else { - std::map> empty_map; - output_single_R(ofs, empty_map, sparse_thr, binary, pv, reduce); + SparseRBlock empty_map; + output_single_R(ofs, empty_map, pv, options); } ++count; } - if (!reduce || GlobalV::DRANK == 0) { + if (!options.reduce || GlobalV::DRANK == 0) + { ofs.close(); } @@ -489,25 +542,13 @@ void ModuleIO::save_sparse( } template void ModuleIO::save_sparse( - const std::map, - std::map>>&, - const std::set>&, - const double&, - const bool&, - const std::string&, + const SparseRMatrix&, + const std::set&, const Parallel_Orbitals&, - const std::string&, - const int&, - const bool&); + const SparseWriteOptions&); template void ModuleIO::save_sparse>( - const std::map, - std::map>>>&, - const std::set>&, - const double&, - const bool&, - const std::string&, + const SparseRMatrix>&, + const std::set&, const Parallel_Orbitals&, - const std::string&, - const int&, - const bool&); + const SparseWriteOptions&); diff --git a/source/source_io/module_hs/write_HS_sparse.h b/source/source_io/module_hs/write_HS_sparse.h index 8791a6bdafd..8640972a847 100644 --- a/source/source_io/module_hs/write_HS_sparse.h +++ b/source/source_io/module_hs/write_HS_sparse.h @@ -4,10 +4,31 @@ #include "source_basis/module_ao/parallel_orbitals.h" #include "source_lcao/LCAO_HS_arrays.hpp" +#include +#include +#include #include namespace ModuleIO { +using RCoordinate = Abfs::Vector3_Order; + +template +using SparseRBlock = std::map>; + +template +using SparseRMatrix = std::map>; + +struct SparseWriteOptions +{ + std::string filename; + std::string label; + double threshold = 0.0; + bool binary = false; + int istep = -1; + bool reduce = true; + std::string temp_dir; +}; void save_dH_sparse(const int& istep, const Parallel_Orbitals& pv, @@ -17,15 +38,10 @@ void save_dH_sparse(const int& istep, const std::string& fileflag = "h"); template -void save_sparse(const std::map, std::map>>& smat, - const std::set>& all_R_coor, - const double& sparse_thr, - const bool& binary, - const std::string& filename, +void save_sparse(const SparseRMatrix& smat, + const std::set& all_R_coor, const Parallel_Orbitals& pv, - const std::string& label, - const int& istep = -1, - const bool& reduce = true); + const SparseWriteOptions& options); } // namespace ModuleIO #endif diff --git a/source/source_io/module_hs/write_vxc_r.hpp b/source/source_io/module_hs/write_vxc_r.hpp index 4fe994d44a2..6445378e37c 100644 --- a/source/source_io/module_hs/write_vxc_r.hpp +++ b/source/source_io/module_hs/write_vxc_r.hpp @@ -57,11 +57,10 @@ void write_Vxc_R(const int nspin, double vtxc = 0.0; // elecstate::PotXC* potxc(&rho_basis, &etxc, vtxc, nullptr); // potxc.cal_v_eff(&chg, &ucell, vr_xc); - elecstate::Potential* potxc - = new elecstate::Potential(&rhod_basis, &rho_basis, &ucell, &vloc, &sf, &solvent, &etxc, &vtxc); + elecstate::Potential potxc(&rhod_basis, &rho_basis, &ucell, &vloc, &sf, &solvent, &etxc, &vtxc); std::vector compnents_list = {"xc"}; - potxc->pot_register(compnents_list); - potxc->update_from_charge(&chg, &ucell); + potxc.pot_register(compnents_list); + potxc.update_from_charge(&chg, &ucell); // 2. allocate H(R) // (the number of hR: 1 for nspin=1, 4; 2 for nspin=2) @@ -89,19 +88,18 @@ void write_Vxc_R(const int nspin, // 3. calculate the Vxc(R) hamilt::HS_Matrix_K vxc_k_ao(pv, 1); // only hk is needed, sk is skipped - std::vector>*> vxcs_op_ao(nspin0); for (int is = 0; is < nspin0; ++is) { - vxcs_op_ao[is] = new hamilt::Veff>(&vxc_k_ao, - kv.kvec_d, - potxc, - &vxcs_R_ao[is], - &ucell, - orb_cutoff, - &gd, - nspin); - vxcs_op_ao[is]->set_current_spin(is); - vxcs_op_ao[is]->contributeHR(); + hamilt::Veff> vxcs_op_ao(&vxc_k_ao, + kv.kvec_d, + &potxc, + &vxcs_R_ao[is], + &ucell, + orb_cutoff, + &gd, + nspin); + vxcs_op_ao.set_current_spin(is); + vxcs_op_ao.contributeHR(); #ifdef __EXX if (GlobalC::exx_info.info_global.cal_exx) { @@ -147,15 +145,18 @@ void write_Vxc_R(const int nspin, { std::set> all_R_coor = sparse_format::get_R_range(vxcs_R_ao[is]); const std::string filename = "Vxc_R_spin" + std::to_string(is); + ModuleIO::SparseWriteOptions options; + options.filename = PARAM.globalv.global_out_dir + filename + ".csr"; + options.label = filename; + options.threshold = sparse_thr; + options.binary = false; + options.istep = -1; + options.reduce = true; + options.temp_dir = PARAM.globalv.global_out_dir; ModuleIO::save_sparse(cal_HR_sparse(vxcs_R_ao[is], sparse_thr), all_R_coor, - sparse_thr, - false, // binary - PARAM.globalv.global_out_dir + filename + ".csr", *pv, - filename, - -1, - true); // all-reduce + options); } } } // namespace ModuleIO diff --git a/source/source_io/module_restart/restart_exx_csr.hpp b/source/source_io/module_restart/restart_exx_csr.hpp index 20fa2bfa1e0..99ff25382ed 100644 --- a/source/source_io/module_restart/restart_exx_csr.hpp +++ b/source/source_io/module_restart/restart_exx_csr.hpp @@ -1,5 +1,6 @@ #pragma once #include "restart_exx_csr.h" +#include "source_base/global_function.h" #include "source_cell/unitcell.h" #include "source_io/module_output/csr_reader.h" #include "source_io/module_hs/write_HS_sparse.h" @@ -114,6 +115,22 @@ namespace ModuleIO ModuleBase::TITLE("ModuleIO", "write_Hexxs_csr"); std::set> all_R_coor; double sparse_threshold = 1e-10; + int matrix_dimension = 0; + for (int it = 0; it < ucell.ntype; ++it) + { + matrix_dimension += ucell.atoms[it].na * ucell.atoms[it].nw * ucell.get_npol(); + } + if (matrix_dimension <= 0) + { + ModuleBase::WARNING_QUIT("ModuleIO::write_Hexxs_csr", + "LCAO matrix dimension must be positive."); + } + Parallel_Orbitals pv; + pv.set_serial(matrix_dimension, matrix_dimension); + const std::string::size_type last_slash = file_name.find_last_of("/\\"); + const std::string temp_dir = last_slash == std::string::npos + ? "" + : file_name.substr(0, last_slash + 1); for (int is = 0;is < Hexxs.size();++is) { for (const auto& HexxA : Hexxs[is]) @@ -126,16 +143,19 @@ namespace ModuleIO all_R_coor.insert(R); } } + ModuleIO::SparseWriteOptions options; + options.filename = file_name + "_" + std::to_string(is) + ".csr"; + options.label = "Hexxs_" + std::to_string(is); + options.threshold = sparse_threshold; + options.binary = false; + options.istep = -1; + options.reduce = false; + options.temp_dir = temp_dir; ModuleIO::save_sparse( calculate_RI_Tensor_sparse(sparse_threshold, Hexxs[is], ucell), all_R_coor, - sparse_threshold, - false, //binary - file_name + "_" + std::to_string(is) + ".csr", - Parallel_Orbitals(), - "Hexxs_" + std::to_string(is), - -1, - false); //no reduce, one file for each process + pv, + options); } } } diff --git a/source/source_io/test/CMakeLists.txt b/source/source_io/test/CMakeLists.txt index e9afa10db23..6dd23247472 100644 --- a/source/source_io/test/CMakeLists.txt +++ b/source/source_io/test/CMakeLists.txt @@ -283,6 +283,23 @@ AddTest( ../../source_basis/module_nao/real_gaunt_table.cpp ../../source_basis/module_nao/radial_collection.cpp ) + +AddTest( + TARGET MODULE_IO_write_hs_r_compat_test + LIBS parameter base ${math_libs} device hcontainer + SOURCES + write_hs_r_compat_test.cpp + ../module_hs/write_HS_R.cpp + ../module_hs/write_HS_sparse.cpp + ../module_hs/single_R_io.cpp + ../module_dm/write_dmr.cpp + ../module_output/ucell_io.cpp + ../module_output/sparse_matrix.cpp + ../module_output/csr_reader.cpp + ../module_output/file_reader.cpp + ../../source_basis/module_ao/parallel_orbitals.cpp + ../../source_lcao/module_hcontainer/test/tmp_mocks.cpp +) endif() AddTest( diff --git a/source/source_io/test/single_R_io_test.cpp b/source/source_io/test/single_R_io_test.cpp index 9454349d0dd..092da754b99 100644 --- a/source/source_io/test/single_R_io_test.cpp +++ b/source/source_io/test/single_R_io_test.cpp @@ -24,6 +24,9 @@ Parallel_Orbitals::~Parallel_Orbitals() void Parallel_2D::set_serial(const int M_A, const int N_A) { + this->nrow = M_A; + this->ncol = N_A; + this->is_serial = true; this->global2local_row_.resize(M_A); this->global2local_row_[0] = 0; this->global2local_row_[1] = 1; @@ -32,6 +35,11 @@ void Parallel_2D::set_serial(const int M_A, const int N_A) this->global2local_row_[4] = -1; //Some rows have global2local_row_ < 0 } +int Parallel_2D::get_global_row_size() const +{ + return this->nrow; +} + TEST(ModuleIOTest, OutputSingleR) { // Create temporary output file @@ -44,16 +52,21 @@ TEST(ModuleIOTest, OutputSingleR) const double sparse_threshold = 1e-8; const bool binary = false; Parallel_Orbitals pv; - PARAM.sys.nlocal = 5; - pv.set_serial(PARAM.sys.nlocal, PARAM.sys.nlocal); + PARAM.sys.nlocal = 99; + pv.set_serial(5, 5); std::map> XR = { {0, {{1, 0.5}, {3, 0.3}}}, {1, {{0, 0.2}, {2, 0.4}}}, {3, {{1, 0.1}, {4, 0.7}}} }; + ModuleIO::SparseWriteOptions options; + options.threshold = sparse_threshold; + options.binary = binary; + options.reduce = true; + options.temp_dir = "./"; // Call function under test - ModuleIO::output_single_R(ofs, XR, sparse_threshold, binary, pv); + ModuleIO::output_single_R(ofs, XR, pv, options); // Close output file and open it for reading ofs.close(); diff --git a/source/source_io/test/write_hs_r_compat_test.cpp b/source/source_io/test/write_hs_r_compat_test.cpp new file mode 100644 index 00000000000..392bafb68d9 --- /dev/null +++ b/source/source_io/test/write_hs_r_compat_test.cpp @@ -0,0 +1,351 @@ +#include "gmock/gmock.h" +#include "gtest/gtest.h" + +#define private public +#include "source_io/module_parameter/parameter.h" +#undef private + +#include "source_base/global_variable.h" +#include "source_cell/module_neighbor/sltk_grid_driver.h" +#include "source_io/module_dm/write_dmr.h" +#include "source_io/module_hs/write_HS_R.h" +#include "source_io/module_hs/write_HS_sparse.h" +#include "source_lcao/module_hcontainer/atom_pair.h" +#include "source_lcao/module_hcontainer/hcontainer.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace sparse_format +{ +void cal_dH(const UnitCell&, + const Parallel_Orbitals&, + LCAO_HS_Arrays&, + const Grid_Driver&, + const TwoCenterBundle&, + const LCAO_Orbitals&, + const int&, + const double&, + const ModuleBase::matrix&) +{ + FAIL() << "cal_dH should not be called by writer compatibility tests."; +} + +void cal_dS(const UnitCell&, + const Parallel_Orbitals&, + LCAO_HS_Arrays&, + const Grid_Driver&, + const TwoCenterBundle&, + const LCAO_Orbitals&, + const double&) +{ + FAIL() << "cal_dS should not be called by writer compatibility tests."; +} + +void cal_TR(const UnitCell&, + const Parallel_Orbitals&, + LCAO_HS_Arrays&, + const Grid_Driver&, + const TwoCenterBundle&, + const LCAO_Orbitals&, + const double&) +{ + FAIL() << "cal_TR should not be called by writer compatibility tests."; +} + +template +void cal_SR(const Parallel_Orbitals&, + std::set>&, + std::map, std::map>>&, + std::map, std::map>>>&, + const Grid_Driver&, + const double&, + hamilt::Hamilt*) +{ + FAIL() << "cal_SR should not be called by writer compatibility tests."; +} + +void destroy_dH_R_sparse(LCAO_HS_Arrays&) {} +void destroy_HS_R_sparse(LCAO_HS_Arrays&) {} +void destroy_T_R_sparse(LCAO_HS_Arrays&) {} + +template void cal_SR( + const Parallel_Orbitals&, + std::set>&, + std::map, std::map>>&, + std::map, std::map>>>&, + const Grid_Driver&, + const double&, + hamilt::Hamilt*); + +template void cal_SR>( + const Parallel_Orbitals&, + std::set>&, + std::map, std::map>>&, + std::map, std::map>>>&, + const Grid_Driver&, + const double&, + hamilt::Hamilt>*); +} // namespace sparse_format + +namespace +{ +std::string read_file(const std::string& filename) +{ + std::ifstream ifs(filename.c_str()); + std::ostringstream oss; + oss << ifs.rdbuf(); + return oss.str(); +} + +std::vector read_binary_ints(const std::string& filename, const size_t count) +{ + std::ifstream ifs(filename.c_str(), std::ios::binary); + std::vector values(count, 0); + for (size_t i = 0; i < count; ++i) + { + ifs.read(reinterpret_cast(&values[i]), sizeof(int)); + } + return values; +} + +int count_substr(const std::string& text, const std::string& pattern) +{ + int count = 0; + std::string::size_type pos = 0; + while ((pos = text.find(pattern, pos)) != std::string::npos) + { + ++count; + pos += pattern.size(); + } + return count; +} + +void init_unitcell(UnitCell& ucell) +{ + ucell.latName = "user_defined_lattice"; + ucell.lat0 = 10.0; + ucell.latvec.e11 = 1.0; + ucell.latvec.e22 = 1.0; + ucell.latvec.e33 = 1.0; + ucell.ntype = 1; + ucell.nat = 1; + ucell.atoms = new Atom[1]; + ucell.set_atom_flag = true; + ucell.atoms[0].label = "Si"; + ucell.atoms[0].na = 1; + ucell.atoms[0].nw = 2; + ucell.atoms[0].taud.resize(1); + ucell.atoms[0].taud[0] = ModuleBase::Vector3(0.0, 0.25, 0.5); +} + +void init_serial_orbitals(Parallel_Orbitals& pv) +{ + pv.atom_begin_row.resize(2); + pv.atom_begin_col.resize(2); + pv.atom_begin_row[0] = 0; + pv.atom_begin_row[1] = 2; + pv.atom_begin_col[0] = 0; + pv.atom_begin_col[1] = 2; + pv.nrow = 2; + pv.ncol = 2; + pv.set_serial(2, 2); +} + +void fill_matrix(hamilt::HContainer& matrix, Parallel_Orbitals& pv, double* values) +{ + hamilt::AtomPair pair(0, 0, 0, 0, 0, &pv, values); + matrix.insert_pair(pair); +} + +bool starts_with(const std::string& text, const std::string& prefix) +{ + return text.find(prefix) == 0; +} +} // namespace + +TEST(WriteHsRCompatibility, FileNameHelpersKeepCurrentContract) +{ + EXPECT_EQ(ModuleIO::hsr_gen_fname("hrs", 0, true, -1), "hrs1_nao.csr"); + EXPECT_EQ(ModuleIO::hsr_gen_fname("hrs", 1, true, 0), "hrs2_nao.csr"); + EXPECT_EQ(ModuleIO::hsr_gen_fname("hrs", 1, false, 3), "hrs2g4_nao.csr"); + EXPECT_EQ(ModuleIO::hsr_gen_fname("srs", 0, false, 0), "srs1g1_nao.csr"); + EXPECT_EQ(ModuleIO::hsr_gen_fname("srs", 0, false, -1), "srs1_nao.csr"); + + EXPECT_EQ(ModuleIO::dhr_gen_fname("dhrx", 0, true, -1), "dhrxrs1_nao.csr"); + EXPECT_EQ(ModuleIO::dhr_gen_fname("dhrx", 0, false, 0), "dhrxrs1g1_nao.csr"); + EXPECT_EQ(ModuleIO::dhr_gen_fname("dsry", 1, false, 2), "dsryrs2g3_nao.csr"); + + EXPECT_EQ(ModuleIO::dmr_gen_fname(1, 0, true, -1), "dmrs1_nao.csr"); + EXPECT_EQ(ModuleIO::dmr_gen_fname(1, 1, false, 2), "dmrs2g3_nao.csr"); + EXPECT_EQ(ModuleIO::dmr_gen_fname(2, 0, true, 5), "dmrs1_nao.npz"); +} + +TEST(WriteHsRCompatibility, HContainerCsrHeaderKeepsCurrentFormat) +{ + const std::string filename = "write_hs_r_header_h.csr"; + std::remove(filename.c_str()); + + UnitCell ucell; + init_unitcell(ucell); + Parallel_Orbitals pv; + init_serial_orbitals(pv); + hamilt::HContainer matrix(&pv); + double values[4] = {1.0, 0.0, 0.5, 2.0}; + fill_matrix(matrix, pv, values); + + ModuleIO::write_hcontainer_csr(filename, &ucell, 5, &matrix, 0, 0, 1, "H"); + + const std::string output = read_file(filename); + EXPECT_THAT(output, testing::HasSubstr(" --- Ionic Step 1 ---\n")); + EXPECT_THAT(output, testing::HasSubstr(" # print H matrix in real space H(R)\n")); + EXPECT_THAT(output, testing::HasSubstr(" 1 # number of spin directions\n")); + EXPECT_THAT(output, testing::HasSubstr(" 1 # spin index\n")); + EXPECT_THAT(output, testing::HasSubstr(" 2 # number of localized basis\n")); + EXPECT_THAT(output, testing::HasSubstr(" 1 # number of Bravais lattice vector R\n")); + EXPECT_THAT(output, testing::HasSubstr(" user_defined_lattice\n")); + EXPECT_THAT(output, testing::HasSubstr(" Si\n")); + EXPECT_THAT(output, testing::HasSubstr(" Direct\n")); + EXPECT_THAT(output, testing::HasSubstr(" 0 0.25 0.5\n")); + EXPECT_THAT(output, testing::HasSubstr(" # CSR Format #\n")); + EXPECT_THAT(output, testing::HasSubstr(" 0 0 0 3\n")); + EXPECT_THAT(output, testing::HasSubstr(" # CSR values\n")); + EXPECT_THAT(output, testing::HasSubstr(" 1.00000e+00 5.00000e-01 2.00000e+00")); + + std::remove(filename.c_str()); +} + +TEST(WriteHsRCompatibility, HContainerCsrAppendKeepsCurrentStepSections) +{ + const std::string filename = "write_hs_r_append_s.csr"; + std::remove(filename.c_str()); + + UnitCell ucell; + init_unitcell(ucell); + Parallel_Orbitals pv; + init_serial_orbitals(pv); + hamilt::HContainer matrix(&pv); + double values[4] = {1.0, 0.0, 0.0, 1.0}; + fill_matrix(matrix, pv, values); + + ModuleIO::write_hcontainer_csr(filename, &ucell, 4, &matrix, 0, 0, 1, "S"); + ModuleIO::write_hcontainer_csr(filename, &ucell, 4, &matrix, 1, 0, 1, "S"); + + const std::string output = read_file(filename); + EXPECT_EQ(count_substr(output, " --- Ionic Step "), 2); + EXPECT_THAT(output, testing::HasSubstr(" --- Ionic Step 1 ---\n")); + EXPECT_THAT(output, testing::HasSubstr(" --- Ionic Step 2 ---\n")); + EXPECT_EQ(count_substr(output, " # print S matrix in real space S(R)\n"), 2); + + std::remove(filename.c_str()); +} + +TEST(WriteHsRCompatibility, DmrCsrHeaderKeepsCurrentFormat) +{ + std::string filename = "write_hs_r_header_dmr.csr"; + std::remove(filename.c_str()); + + UnitCell ucell; + init_unitcell(ucell); + Parallel_Orbitals pv; + init_serial_orbitals(pv); + hamilt::HContainer matrix(&pv); + double values[4] = {0.25, 0.0, 0.0, 0.75}; + fill_matrix(matrix, pv, values); + + ModuleIO::write_dmr_csr(filename, &ucell, 4, &matrix, 0, 1, 2); + + const std::string output = read_file(filename); + EXPECT_THAT(output, testing::HasSubstr(" --- Ionic Step 1 ---\n")); + EXPECT_THAT(output, testing::HasSubstr(" # print density matrix in real space DM(R)\n")); + EXPECT_THAT(output, testing::HasSubstr(" 2 # number of spin directions\n")); + EXPECT_THAT(output, testing::HasSubstr(" 2 # spin index\n")); + EXPECT_THAT(output, testing::HasSubstr(" 2 # number of localized basis\n")); + EXPECT_THAT(output, testing::HasSubstr(" 1 # number of Bravais lattice vector R\n")); + + std::remove(filename.c_str()); +} + +TEST(WriteHsRCompatibility, LegacySparseHeaderKeepsStepStyle) +{ + const std::string filename = "write_hs_r_legacy_s.csr"; + std::remove(filename.c_str()); + + GlobalV::DRANK = 0; + PARAM.sys.global_out_dir = "./"; + PARAM.sys.nlocal = 99; + + Parallel_Orbitals pv; + init_serial_orbitals(pv); + const Abfs::Vector3_Order r_vector(0, 0, 0); + std::set> all_R_coor; + all_R_coor.insert(r_vector); + std::map, std::map>> sparse_matrix; + sparse_matrix[r_vector][0][1] = 0.5; + sparse_matrix[r_vector][1][1] = 1.5; + + ModuleIO::SparseWriteOptions options; + options.filename = filename; + options.label = "S"; + options.threshold = 1e-10; + options.binary = false; + options.istep = 0; + options.reduce = false; + options.temp_dir = "./"; + ModuleIO::save_sparse(sparse_matrix, all_R_coor, pv, options); + + const std::string output = read_file(filename); + EXPECT_TRUE(starts_with(output, "STEP: 0\n")); + EXPECT_THAT(output, testing::HasSubstr("Matrix Dimension of S(R): 2\n")); + EXPECT_THAT(output, testing::HasSubstr("Matrix number of S(R): 1\n")); + EXPECT_THAT(output, testing::HasSubstr("0 0 0 2\n")); + + std::remove(filename.c_str()); +} + +TEST(WriteHsRCompatibility, LegacySparseBinaryHeaderWritesConcreteStep) +{ + const std::string filename = "write_hs_r_legacy_binary_s.csr"; + std::remove(filename.c_str()); + + GlobalV::DRANK = 0; + PARAM.sys.global_out_dir = "./"; + PARAM.sys.nlocal = 99; + + Parallel_Orbitals pv; + init_serial_orbitals(pv); + const Abfs::Vector3_Order r_vector(0, 0, 0); + std::set> all_R_coor; + all_R_coor.insert(r_vector); + std::map, std::map>> sparse_matrix; + sparse_matrix[r_vector][0][1] = 0.5; + sparse_matrix[r_vector][1][1] = 1.5; + + ModuleIO::SparseWriteOptions options; + options.filename = filename; + options.label = "S"; + options.threshold = 1e-10; + options.binary = true; + options.istep = 3; + options.reduce = false; + options.temp_dir = "./"; + ModuleIO::save_sparse(sparse_matrix, all_R_coor, pv, options); + + const std::vector header_and_r = read_binary_ints(filename, 7); + EXPECT_THAT(header_and_r, testing::ElementsAre(3, 2, 1, 0, 0, 0, 2)); + + std::remove(filename.c_str()); +} + +TEST(WriteHsRCompatibility, HeaderStyleSamplesRemainDistinct) +{ + EXPECT_TRUE(starts_with(" --- Ionic Step 1 ---", " --- Ionic Step")); + EXPECT_TRUE(starts_with("STEP: 0", "STEP:")); + EXPECT_TRUE(starts_with("IONIC_STEP: 1", "IONIC_STEP:")); +}