Feature: Add a output interface for hcontainer and a sparse matrix cl…

…ass for CSR io (#2745)

source/module_hamilt_lcao/module_hcontainer/CMakeLists.txt

@@ -2,6 +2,7 @@ list(APPEND objects
     base_matrix.cpp
     atom_pair.cpp
     hcontainer.cpp
+    output_hcontainer.cpp
 )
 
 add_library(

source/module_hamilt_lcao/module_hcontainer/output_hcontainer.cpp

@@ -0,0 +1,114 @@
+#include "output_hcontainer.h"
+
+#include <fstream>
+
+#include "module_io/sparse_matrix.h"
+
+namespace hamilt
+{
+
+/**
+ * @brief Constructor of Output_HContainer
+ * @attention ofs should be open outside of this interface
+ */
+template <typename T>
+Output_HContainer<T>::Output_HContainer(hamilt::HContainer<T>* hcontainer,
+                                        const Parallel_Orbitals* ParaV,
+                                        const UnitCell& ucell,
+                                        std::ostream& ofs,
+                                        double sparse_threshold,
+                                        int precision)
+    : _hcontainer(hcontainer),
+      _ParaV(ParaV),
+      _ucell(ucell),
+      _ofs(ofs),
+      _sparse_threshold(sparse_threshold),
+      _precision(precision)
+{
+}
+
+template <typename T>
+void Output_HContainer<T>::write()
+{
+    int size_for_loop_R = this->_hcontainer->size_R_loop();
+    int rx, ry, rz;
+    for (int iR = 0; iR < size_for_loop_R; iR++)
+    {
+        this->_hcontainer->loop_R(iR, rx, ry, rz);
+        this->write_single_R(rx, ry, rz);
+    }
+}
+
+template <typename T>
+void Output_HContainer<T>::write(int rx_in, int ry_in, int rz_in)
+{
+    int size_for_loop_R = this->_hcontainer->size_R_loop();
+    int rx, ry, rz;
+    int find_R = 0;
+    for (int iR = 0; iR < size_for_loop_R; iR++)
+    {
+        this->_hcontainer->loop_R(iR, rx, ry, rz);
+        if (rx == rx_in && ry == ry_in && rz == rz_in)
+        {
+            find_R += 1;
+            this->write_single_R(rx, ry, rz);
+            break;
+        }
+    }
+    if (find_R == 0)
+    {
+        ModuleBase::WARNING_QUIT("Output_HContainer::write", "Cannot find the R vector from the HContaine");
+    }
+}
+
+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>(this->_ParaV->nrow, this->_ParaV->ncol);
+    int nonzero = 0;
+    for (int iap = 0; iap < this->_hcontainer->size_atom_pairs(); ++iap)
+    {
+        auto atom_pair = this->_hcontainer->get_atom_pair(iap);
+        int iat1 = atom_pair.get_atom_i();
+        int iat2 = atom_pair.get_atom_j();
+        int T1 = _ucell.iat2it[iat1];
+        int T2 = _ucell.iat2it[iat2];
+        int I1 = _ucell.iat2ia[iat1];
+        int I2 = _ucell.iat2ia[iat2];
+        const int start1 = _ucell.itiaiw2iwt(T1, I1, 0);
+        const int start2 = _ucell.itiaiw2iwt(T2, I2, 0);
+        int size1 = _ucell.atoms[T1].nw;
+        int size2 = _ucell.atoms[T2].nw;
+        for (int iw1 = 0; iw1 < size1; ++iw1)
+        {
+            const int global_index1 = start1 + iw1;
+            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
+                }
+            }
+        }
+    }
+    if (nonzero != 0)
+    {
+        _ofs << rx << " " << ry << " " << rz << " " << nonzero << std::endl;
+        sparse_matrix.printToCSR(_ofs, _sparse_threshold, _precision);
+    }
+    this->_hcontainer->unfix_R();
+}
+
+// explicit instantiation of template class with double type
+template class Output_HContainer<double>;
+// to do: explicit instantiation of template class with std::complex<double> type
+// template class Output_HContainer<std::complex<double>>;
+
+} // namespace hamilt

source/module_hamilt_lcao/module_hcontainer/output_hcontainer.h

@@ -0,0 +1,49 @@
+#ifndef OUTPUT_HCONTAINER_H
+#define OUTPUT_HCONTAINER_H
+
+#include "module_cell/unitcell.h"
+#include "module_hamilt_lcao/module_hcontainer/hcontainer.h"
+
+namespace hamilt
+{
+
+/**
+ * @brief A class to output the HContainer
+ */
+template <typename T>
+class Output_HContainer
+{
+  public:
+    Output_HContainer(hamilt::HContainer<T>* hcontainer,
+                      const Parallel_Orbitals* ParaV,
+                      const UnitCell& ucell,
+                      std::ostream& ofs,
+                      double sparse_threshold,
+                      int precision);
+    // write the matrices of all R vectors to the output stream
+    void write();
+
+    /**
+     * write the matrix of a single R vector to the output stream
+     * rx_in, ry_in, rz_in: the R vector from the input
+     */
+    void write(int rx_in, int ry_in, int rz_in);
+
+    /**
+     * write the matrix of a single R vector to the output stream
+     * rx, ry, rz: the R vector from the HContainer
+     */
+    void write_single_R(int rx, int ry, int rz);
+
+  private:
+    hamilt::HContainer<T>* _hcontainer;
+    const UnitCell& _ucell;
+    const Parallel_Orbitals* _ParaV;
+    std::ostream& _ofs;
+    double _sparse_threshold;
+    int _precision;
+};
+
+} // namespace ModuleIO
+
+#endif // OUTPUT_HCONTAINER_H

source/module_hamilt_lcao/module_hcontainer/test/CMakeLists.txt

@@ -20,4 +20,10 @@ AddTest(
   SOURCES test_hcontainer_time.cpp ../base_matrix.cpp ../hcontainer.cpp ../atom_pair.cpp  
   ../../../module_basis/module_ao/parallel_2d.cpp ../../../module_basis/module_ao/parallel_orbitals.cpp tmp_mocks.cpp
 )
+
+AddTest(
+  TARGET test_output_hcontainer
+  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
+)
 endif() 

source/module_hamilt_lcao/module_hcontainer/test/output_hcontainer_test.cpp

@@ -0,0 +1,164 @@
+#include "../output_hcontainer.h"
+
+#include "../hcontainer.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "module_cell/unitcell.h"
+
+/************************************************
+ *  unit test of output_hcontainer.cpp
+ ***********************************************/
+
+/**
+ * - Tested Functions:
+ * - write()
+ * - write the matrices of all R vectors to the output stream
+ * - write(int rx_in, int ry_in, int rz_in)
+ * - write the matrix of a single R vector to the output stream
+ * - write_single_R(int rx, int ry, int rz)
+ * - write the matrix of a single R vector to the output stream
+ */
+
+class OutputHContainerTest : public testing::Test
+{
+  protected:
+    UnitCell ucell;
+    std::string output;
+    void SetUp() override
+    {
+        ucell.ntype = 1;
+        ucell.nat = 2;
+        ucell.atoms = new Atom[ucell.ntype];
+        ucell.iat2it = new int[ucell.nat];
+        ucell.iat2ia = new int[ucell.nat];
+        for (int iat = 0; iat < ucell.nat; iat++)
+        {
+            ucell.iat2ia[iat] = iat;
+            ucell.iat2it[iat] = 0;
+        }
+        ucell.atoms[0].nw = 2;
+        ucell.iwt2iat = new int[4];
+        ucell.iwt2iw = new int[4];
+        ucell.itia2iat.create(ucell.ntype, ucell.nat);
+        ucell.iat2iwt.resize(ucell.nat);
+        ucell.iat2iwt[0] = 0;
+        ucell.iat2iwt[1] = 2;
+        ucell.itia2iat(0, 0) = 0;
+        ucell.itia2iat(0, 1) = 1;
+        ucell.iwt2iat[0] = 0;
+        ucell.iwt2iat[1] = 0;
+        ucell.iwt2iat[2] = 1;
+        ucell.iwt2iat[3] = 1;
+        ucell.iwt2iw[0] = 0;
+        ucell.iwt2iw[1] = 1;
+        ucell.iwt2iw[2] = 0;
+        ucell.iwt2iw[3] = 1;
+    }
+    void TearDown() override
+    {
+        delete[] ucell.atoms;
+        delete[] ucell.iat2it;
+        delete[] ucell.iat2ia;
+        delete[] ucell.iwt2iat;
+        delete[] ucell.iwt2iw;
+    }
+};
+
+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[i] = i * 2;
+        ParaV.atom_begin_col[i] = i * 2;
+    }
+    ParaV.nrow = 4;
+    ParaV.ncol = 4;
+    std::ofstream ofs("output_hcontainer.log");
+    ParaV.set_global2local(4, 4, false, ofs);
+    // 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) << " "
+    //           << ParaV.global2local_col(2) << " " << ParaV.global2local_col(3) << std::endl;
+    ofs.close();
+    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};
+    // 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);
+    HR.insert_pair(ap1);
+    HR.insert_pair(ap2);
+    /*
+    for (int ir = 0; ir < HR.size_R_loop(); ++ir)
+    {
+        int rx, ry, rz;
+        HR.loop_R(ir, rx, ry, rz);
+        HR.fix_R(rx, ry, rz);
+        // std::cout << "rx = " << rx << " ry = " << ry << " rz = " << rz << std::endl;
+        for (int iap = 0; iap < HR.size_atom_pairs(); ++iap)
+        {
+            hamilt::AtomPair<double>& tmp_ap = HR.get_atom_pair(iap);
+            if (rx == 0 && ry == 1 && rz == 1)
+            {
+                EXPECT_DOUBLE_EQ(tmp_ap.get_value(0, 0), 0);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_value(0, 1), 4);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_value(1, 0), 7);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_value(1, 1), 0);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_matrix_value(0, 2), 0);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_matrix_value(0, 3), 4);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_matrix_value(1, 2), 7);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_matrix_value(1, 3), 0);
+            }
+            else if (rx == 0 && ry == 0 && rz == 0)
+            {
+                EXPECT_DOUBLE_EQ(tmp_ap.get_value(0, 0), 5);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_value(0, 1), 6);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_value(1, 0), 0);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_value(1, 1), 10);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_matrix_value(2, 2), 5);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_matrix_value(2, 3), 6);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_matrix_value(3, 2), 0);
+                EXPECT_DOUBLE_EQ(tmp_ap.get_matrix_value(3, 3), 10);
+            }
+        }
+        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
+    testing::internal::CaptureStdout();
+    output_HR.write(0, 1, 1);
+    output = testing::internal::GetCapturedStdout();
+    EXPECT_THAT(output, testing::HasSubstr("0 1 1 2"));
+    EXPECT_THAT(output, testing::HasSubstr(" 4.00e+00 7.00e+00"));
+    EXPECT_THAT(output, testing::HasSubstr(" 3 2"));
+    EXPECT_THAT(output, testing::HasSubstr(" 0 1 2 2 2"));
+    // the second R
+    testing::internal::CaptureStdout();
+    output_HR.write(0, 0, 0);
+    output = testing::internal::GetCapturedStdout();
+    EXPECT_THAT(output, testing::HasSubstr("0 0 0 3"));
+    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"));
+    // output all R
+    testing::internal::CaptureStdout();
+    output_HR.write();
+    output = testing::internal::GetCapturedStdout();
+    EXPECT_THAT(output, testing::HasSubstr("0 1 1 2"));
+    EXPECT_THAT(output, testing::HasSubstr(" 4.00e+00 7.00e+00"));
+    EXPECT_THAT(output, testing::HasSubstr(" 3 2"));
+    EXPECT_THAT(output, testing::HasSubstr(" 0 1 2 2 2"));
+    EXPECT_THAT(output, testing::HasSubstr("0 0 0 3"));
+    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"));
+}

source/module_io/CMakeLists.txt

@@ -50,6 +50,7 @@ if(ENABLE_LCAO)
       dos_nao.cpp
       output_dm.cpp
       output_dm1.cpp
+      sparse_matrix.cpp
   )
   list(APPEND objects_advanced
       unk_overlap_lcao.cpp