Feature: Add CSR file reader classes in module_io

source/module_hamilt_lcao/module_hcontainer/output_hcontainer.cpp

@@ -66,7 +66,7 @@ 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>(this->_ParaV->nrow, this->_ParaV->ncol);
-    int nonzero = 0;
+    sparse_matrix.setSparseThreshold(this->_sparse_threshold);
     for (int iap = 0; iap < this->_hcontainer->size_atom_pairs(); ++iap)
     {
         auto atom_pair = this->_hcontainer->get_atom_pair(iap);
@@ -86,22 +86,16 @@ void Output_HContainer<T>::write_single_R(int rx, int ry, int rz)
             for (int iw2 = 0; iw2 < size2; ++iw2)
             {
                 const int global_index2 = start2 + iw2;
-
                 T tmp_matrix_value = atom_pair.get_matrix_value(global_index1, global_index2);
-
-                if (std::abs(tmp_matrix_value) > _sparse_threshold)
-                {
-                    nonzero++;
-                    sparse_matrix.addValue(global_index1, global_index2, tmp_matrix_value);
-                    // to do: consider 2D block-cyclic distribution
-                }
+                sparse_matrix.insert(global_index1, global_index2, tmp_matrix_value);
+                // to do: consider 2D block-cyclic distribution
             }
         }
     }
-    if (nonzero != 0)
+    if (sparse_matrix.getNNZ() != 0)
     {
-        _ofs << rx << " " << ry << " " << rz << " " << nonzero << std::endl;
-        sparse_matrix.printToCSR(_ofs, _sparse_threshold, _precision);
+        _ofs << rx << " " << ry << " " << rz << " " << sparse_matrix.getNNZ() << std::endl;
+        sparse_matrix.printToCSR(_ofs, _precision);
     }
     this->_hcontainer->unfix_R();
 }

source/module_hamilt_lcao/module_hcontainer/test/CMakeLists.txt

@@ -22,8 +22,16 @@ AddTest(
 )
 
 AddTest(
-  TARGET test_output_hcontainer
+  TARGET test_hcontainer_output
   LIBS base ${math_libs} device
-  SOURCES output_hcontainer_test.cpp ../output_hcontainer.cpp ../../../module_basis/module_ao/parallel_2d.cpp ../../../module_basis/module_ao/parallel_orbitals.cpp ../../../module_io/sparse_matrix.cpp ../base_matrix.cpp ../hcontainer.cpp ../atom_pair.cpp tmp_mocks.cpp
+  SOURCES output_hcontainer_test.cpp
+    tmp_mocks.cpp
+    ../output_hcontainer.cpp
+    ../base_matrix.cpp
+    ../hcontainer.cpp
+    ../atom_pair.cpp
+    ../../../module_basis/module_ao/parallel_2d.cpp
+    ../../../module_basis/module_ao/parallel_orbitals.cpp
+    ../../../module_io/sparse_matrix.cpp
 )
 endif() 

source/module_hamilt_lcao/module_hcontainer/test/output_hcontainer_test.cpp

@@ -67,9 +67,16 @@ class OutputHContainerTest : public testing::Test
 TEST_F(OutputHContainerTest, Write)
 {
     Parallel_Orbitals ParaV;
-    ParaV.atom_begin_row.resize(2);
-    ParaV.atom_begin_col.resize(2);
-    for (int i = 0; i < 2; i++)
+    ParaV.atom_begin_row.resize(3);
+    ParaV.atom_begin_col.resize(3);
+    // Here is a bug, "nat" is 2 in this test, but
+    // the Parallel_Orbitals::get_col_size function
+    // will run "iat += 1" even if iat == nat - 1
+    // resulting in wrong indexing of atom_begin_col
+    // so it is resized to 3 to avoid this bug temporarily
+    // @ 2023-07-26
+    // this bug is expected to be fixed soon
+    for (int i = 0; i < 3; i++)
     {
         ParaV.atom_begin_row[i] = i * 2;
         ParaV.atom_begin_col[i] = i * 2;
@@ -86,16 +93,16 @@ TEST_F(OutputHContainerTest, Write)
     remove("output_hcontainer.log");
     hamilt::HContainer<double> HR(&ParaV);
     double correct_array[16] = {1, 2, 0, 4, 5, 0, 7, 0, 3, 0, 5, 6, 7, 8, 0, 10};
+    double correct_array1[16] = {1, 2, 0, 4, 5, 0, 7, 0, 3, 0, 5, 6, 7, 8, 0, 10};
     // correct_array represent a matrix of
     // 1 2 0 4
     // 5 0 7 0
     // 3 0 5 6
     // 7 8 0 10
     hamilt::AtomPair<double> ap1(0, 1, 0, 1, 1, &ParaV, correct_array);
-    hamilt::AtomPair<double> ap2(1, 1, 0, 0, 0, &ParaV, correct_array);
+    hamilt::AtomPair<double> ap2(1, 1, 0, 0, 0, &ParaV, correct_array1);
     HR.insert_pair(ap1);
     HR.insert_pair(ap2);
-    /*
     for (int ir = 0; ir < HR.size_R_loop(); ++ir)
     {
         int rx, ry, rz;
@@ -130,7 +137,6 @@ TEST_F(OutputHContainerTest, Write)
         }
         HR.unfix_R();
     }
-    */
     double sparse_threshold = 0.1;
     hamilt::Output_HContainer<double> output_HR(&HR, &ParaV, ucell, std::cout, sparse_threshold, 2);
     // the first R

source/module_io/CMakeLists.txt

@@ -51,6 +51,8 @@ if(ENABLE_LCAO)
       output_dm.cpp
       output_dm1.cpp
       sparse_matrix.cpp
+      file_reader.cpp
+      csr_reader.cpp
   )
   list(APPEND objects_advanced
       unk_overlap_lcao.cpp

source/module_io/csr_reader.cpp

@@ -0,0 +1,140 @@
+#include "csr_reader.h"
+
+#include "module_base/tool_quit.h"
+
+namespace ModuleIO
+{
+
+// constructor
+template <typename T>
+csrFileReader<T>::csrFileReader(const std::string& filename) : FileReader(filename)
+{
+    parseFile();
+}
+
+// function to parse file
+template <typename T>
+void csrFileReader<T>::parseFile()
+{
+    // Check if file is open
+    if (!isOpen())
+    {
+        ModuleBase::WARNING_QUIT("csrFileReader::parseFile", "File is not open");
+    }
+
+    std::string tmp_string;
+
+    // Read the step
+    readLine();
+    ss >> tmp_string >> step;
+
+    // Read the matrix dimension
+    readLine();
+    ss >> tmp_string >> tmp_string >> tmp_string >> tmp_string >> matrixDimension;
+
+    // Read the number of R
+    readLine();
+    ss >> tmp_string >> tmp_string >> tmp_string >> tmp_string >> numberOfR;
+
+    // Read the matrices
+    for (int i = 0; i < numberOfR; i++)
+    {
+        std::vector<int> RCoord(3);
+        int nonZero;
+
+        readLine();
+        ss >> RCoord[0] >> RCoord[1] >> RCoord[2] >> nonZero;
+        RCoordinates.push_back(RCoord);
+
+        std::vector<T> csr_values(nonZero);
+        std::vector<int> csr_col_ind(nonZero);
+        std::vector<int> csr_row_ptr(matrixDimension + 1);
+
+        // read CSR values
+        readLine();
+        for (int i = 0; i < nonZero; i++)
+        {
+            ss >> csr_values[i];
+        }
+        // read column indices
+        readLine();
+        for (int i = 0; i < nonZero; i++)
+        {
+            ss >> csr_col_ind[i];
+        }
+        // read row pointers
+        readLine();
+        for (int i = 0; i < matrixDimension + 1; i++)
+        {
+            ss >> csr_row_ptr[i];
+        }
+
+        SparseMatrix<T> matrix(matrixDimension, matrixDimension);
+        matrix.readCSR(csr_values, csr_col_ind, csr_row_ptr);
+        sparse_matrices.push_back(matrix);
+    }
+}
+
+// function to get R coordinate
+template <typename T>
+std::vector<int> csrFileReader<T>::getRCoordinate(int index) const
+{
+    if (index < 0 || index >= RCoordinates.size())
+    {
+        ModuleBase::WARNING_QUIT("csrFileReader::getRCoordinate", "Index out of range");
+    }
+    return RCoordinates[index];
+}
+
+// function to get matrix
+template <typename T>
+SparseMatrix<T> csrFileReader<T>::getMatrix(int index) const
+{
+    if (index < 0 || index >= sparse_matrices.size())
+    {
+        ModuleBase::WARNING_QUIT("csrFileReader::getMatrix", "Index out of range");
+    }
+    return sparse_matrices[index];
+}
+
+// function to get matrix using R coordinate
+template <typename T>
+SparseMatrix<T> csrFileReader<T>::getMatrix(int Rx, int Ry, int Rz)
+{
+    for (int i = 0; i < RCoordinates.size(); i++)
+    {
+        if (RCoordinates[i][0] == Rx && RCoordinates[i][1] == Ry && RCoordinates[i][2] == Rz)
+        {
+            return sparse_matrices[i];
+        }
+    }
+    ModuleBase::WARNING_QUIT("csrFileReader::getMatrix", "R coordinate not found");
+}
+
+// function to get matrix
+template <typename T>
+int csrFileReader<T>::getNumberOfR() const
+{
+    return numberOfR;
+}
+
+// function to get matrixDimension
+template <typename T>
+int csrFileReader<T>::getMatrixDimension() const
+{
+    return matrixDimension;
+}
+
+// function to get step
+template <typename T>
+int csrFileReader<T>::getStep() const
+{
+    return step;
+}
+
+// T of AtomPair can be double
+template class csrFileReader<double>;
+// ToDo: T of AtomPair can be std::complex<double>
+// template class csrFileReader<std::complex<double>>;
+
+} // namespace ModuleIO

source/module_io/csr_reader.h

@@ -0,0 +1,71 @@
+#ifndef CSR_READER_H
+#define CSR_READER_H
+
+#include <fstream>
+
+#include "file_reader.h"
+#include "sparse_matrix.h"
+
+namespace ModuleIO
+{
+
+/**
+ * @brief Class to read CSR file
+ * @details This class is used to read CSR file
+ * the default format is like:
+ * ```
+ * STEP: 0
+ * Matrix Dimension of S(R): 4
+ * Matrix number of S(R): 2
+ * 0 1 1 2            # (0,1,1) is the R coordinate, 2 is the number of non-zero elements
+ *  4.00e+00 7.00e+00 # non-zero elements
+ *  3 2               # column indices
+ *  0 1 2 2 2         # row pointer
+ * 0 0 0 3            # the second R coordinate and number of non-zero elements
+ * 5.00e+00 6.00e+00 1.00e+01
+ * 2 3 3
+ * 0 0 0 2 3
+ * ```
+ * It will store the R coordinates and sparse matrices as two vectors.
+ * One can use getter functions to get the R coordinates and sparse matrices,
+ * and related info including step, matrix dimension, number of R.
+ */
+template <typename T>
+class csrFileReader : public FileReader
+{
+  public:
+    // Constructor
+    csrFileReader(const std::string& filename);
+
+    // read all matrices of all R coordinates
+    void parseFile();
+
+    // get number of R
+    int getNumberOfR() const;
+
+    // get sparse matrix of a specific R coordinate
+    SparseMatrix<T> getMatrix(int Rx, int Ry, int Rz);
+
+    // get matrix by using index
+    SparseMatrix<T> getMatrix(int index) const;
+
+    // get R coordinate using index
+    std::vector<int> getRCoordinate(int index) const;
+
+    // get step
+    int getStep() const;
+
+    // get matrix dimension
+    int getMatrixDimension() const;
+
+  private:
+    std::vector<std::vector<int>> RCoordinates;
+    std::vector<SparseMatrix<T>> sparse_matrices;
+    int step;
+    int matrixDimension;
+    int numberOfR;
+};
+
+} // namespace ModuleIO
+
+#endif // READ_CSR_H

source/module_io/file_reader.cpp

@@ -0,0 +1,50 @@
+#include "file_reader.h"
+
+#include "module_base/tool_quit.h"
+
+namespace ModuleIO
+{
+
+// Constructor
+FileReader::FileReader(std::string filename)
+{
+    ifs.open(filename.c_str());
+    if (!ifs.is_open())
+    {
+        ModuleBase::WARNING_QUIT("FileReader::FileReader", "Error opening file");
+    }
+}
+
+// Destructor
+FileReader::~FileReader()
+{
+    if (ifs.is_open())
+    {
+        ifs.close();
+    }
+}
+
+// Function to check if file is open
+bool FileReader::isOpen() const
+{
+    return ifs.is_open();
+}
+
+// Function to read a line and return string stream
+void FileReader::readLine()
+{
+    // add warning if file is not open
+    if (!ifs.eof())
+    {
+        std::string line;
+        std::getline(ifs, line);
+        ss.clear();
+        ss.str(line);
+    }
+    else
+    {
+        ModuleBase::WARNING_QUIT("FileReader::readLine", "End of file");
+    }
+}
+
+} // namespace ModuleIO