Perf&Refactor: optimize the performanace of psi_initializer with omp

source/Makefile.Objects

@@ -523,6 +523,7 @@ OBJS_IO=input_conv.o\
     read_input_item_other.o\
     read_input_item_output.o\
     read_set_globalv.o\
+    orb_io.o\
 
 OBJS_IO_LCAO=cal_r_overlap_R.o\
       write_orb_info.o\

source/module_basis/module_nao/atomic_radials.cpp

@@ -3,11 +3,16 @@
 #include "module_base/math_integral.h"
 #include "module_base/parallel_common.h"
 #include "module_base/tool_quit.h"
+
+// FIXME: should update with pyabacus
+// #include "module_io/orb_io.h"
+
 #include "projgen.h"
 
 #include <fstream>
 #include <iostream>
 #include <string>
+#include <numeric>
 
 AtomicRadials& AtomicRadials::operator=(const AtomicRadials& rhs)
 {
@@ -207,7 +212,6 @@ void AtomicRadials::read_abacus_orb(std::ifstream& ifs, std::ofstream* ptr_log,
     {
         rgrid[ir] = ir * dr;
     }
-
     chi_ = new NumericalRadial[nchi_];
 
     // record whether an orbital has been read or not
@@ -269,116 +273,47 @@ void AtomicRadials::read_abacus_orb(std::ifstream& ifs, std::ofstream* ptr_log,
     delete[] rgrid;
 }
 
-void AtomicRadials::read_abacus_orb(std::ifstream& ifs,
-                                    std::string& elem,
-                                    double& ecut,
-                                    int& nr,
-                                    double& dr,
-                                    std::vector<int>& nzeta,
-                                    std::vector<std::vector<double>>& radials,
-                                    const int rank)
-{
-    nr = 0; // number of grid points
-    dr = 0; // grid spacing
-    int lmax, nchi = 0; // number of radial functions
-    std::vector<std::vector<int>> radial_map_; // build a map from [l][izeta] to 1-d array index
-    std::string tmp;
-    // first read the header
-    if (rank == 0)
-    {
-        if (!ifs.is_open())
-        {
-            ModuleBase::WARNING_QUIT("AtomicRadials::read_abacus_orb", "Couldn't open orbital file.");
-        }
-        while (ifs >> tmp)
-        {
-            if (tmp == "Element")
-            {
-                ifs >> elem;
-            }
-            else if (tmp == "Cutoff(Ry)")
-            {
-                ifs >> ecut;
-            }
-            else if (tmp == "Lmax")
-            {
-                ifs >> lmax;
-                nzeta.resize(lmax + 1);
-                for (int l = 0; l <= lmax; ++l)
-                {
-                    ifs >> tmp >> tmp >> tmp >> nzeta[l];
-                }
-            }
-            else if (tmp == "Mesh")
-            {
-                ifs >> nr;
-                continue;
-            }
-            else if (tmp == "dr")
-            {
-                ifs >> dr;
-                break;
-            }
-        }
-        radial_map_.resize(lmax + 1);
-        for (int l = 0; l <= lmax; ++l)
-        {
-            radial_map_[l].resize(nzeta[l]);
-        }
-        int ichi = 0;
-        for (int l = 0; l <= lmax; ++l)
-        {
-            for (int iz = 0; iz < nzeta[l]; ++iz)
-            {
-                radial_map_[l][iz] = ichi++; // return the value of ichi, then increment
-            }
-        }
-        nchi = ichi; // total number of radial functions
-        radials.resize(nchi);
-        std::for_each(radials.begin(), radials.end(), [nr](std::vector<double>& v) { v.resize(nr); });
-    }
-
-    // broadcast the header information
-#ifdef __MPI
-    Parallel_Common::bcast_string(elem);
-    Parallel_Common::bcast_double(ecut);
-    Parallel_Common::bcast_int(lmax);
-    Parallel_Common::bcast_int(nchi);
-    Parallel_Common::bcast_int(nr);
-    Parallel_Common::bcast_double(dr);
-#endif
-
-    // then adjust the size of the vectors
-    if (rank != 0)
-    {
-        nzeta.resize(lmax + 1);
-        radials.resize(nchi);
-        std::for_each(radials.begin(), radials.end(), [nr](std::vector<double>& v) { v.resize(nr); });
-    }
-    // broadcast the number of zeta functions for each angular momentum
-#ifdef __MPI
-    Parallel_Common::bcast_int(nzeta.data(), lmax + 1);
-#endif
-
-    // read the radial functions by rank0
-    int ichi = 0;
-    for (int i = 0; i != nchi; ++i)
-    {
-        if (rank == 0)
-        {
-            int l, izeta;
-            ifs >> tmp >> tmp >> tmp;
-            ifs >> tmp >> l >> izeta;
-            ichi = radial_map_[l][izeta];
-            for (int ir = 0; ir != nr; ++ir)
-            {
-                ifs >> radials[ichi][ir];
-            }
-        }
-    // broadcast the radial functions
-#ifdef __MPI
-        Parallel_Common::bcast_int(ichi); // let other ranks know where to store the radial function
-        Parallel_Common::bcast_double(radials[ichi].data(), nr);
-#endif
-    }
-}
+// FIXME: should update with pyabacus
+// void AtomicRadials::read_abacus_orb(std::ifstream& ifs, std::ofstream* ptr_log, const int rank)
+// {
+//     /*
+//      * Read the orbital file.
+//      *
+//      * For orbital file format, see
+//      * (new) abacus-develop/tools/SIAB/PyTorchGradient/source/IO/print_orbital.py
+//      * (old) abacus-develop/tools/SIAB/SimulatedAnnealing/source/src_spillage/Plot_Psi.cpp
+//      *                                                                                  */
+//     int ngrid = 0; // number of grid points
+//     double dr = 0; // grid spacing
+//     std::string tmp;
+
+//     int nr;
+//     std::vector<int> nzeta;
+//     std::vector<std::vector<double>> radials;
+
+//     ModuleIO::read_abacus_orb(ifs, symbol_, orb_ecut_, nr, dr, nzeta, radials, rank);
+
+//     lmax_ = nzeta.size() - 1;
+//     nzeta_ = new int[lmax_ + 1];
+//     std::copy(nzeta.begin(), nzeta.end(), nzeta_);
+
+//     nchi_ = std::accumulate(nzeta.begin(), nzeta.end(), 0);
+//     nzeta_max_ = *std::max_element(nzeta.begin(), nzeta.end());
+
+//     indexing();
+
+//     std::vector<double> rgrid(nr);
+//     std::iota(rgrid.begin(), rgrid.end(), 0);
+//     std::for_each(rgrid.begin(), rgrid.end(), [dr](double& r) { r *= dr; });
+//     chi_ = new NumericalRadial[nchi_];
+//     int ichi = 0;
+//     for (int l = 0; l <= lmax_; ++l)
+//     {
+//         for (int izeta = 0; izeta < nzeta[l]; ++izeta)
+//         {
+//             chi_[index(l, izeta)].build(l, true, nr, rgrid.data(), radials[ichi].data(), 0, izeta, symbol_, itype_, false);
+//             chi_[index(l, izeta)].normalize();
+//             ++ichi;
+//         }
+//     }
+// }

source/module_basis/module_nao/atomic_radials.h

@@ -43,28 +43,6 @@ class AtomicRadials : public RadialSet
     //! Get the energy cutoff as given by the orbital file
     double orb_ecut() const { return orb_ecut_; }
 
-    /**
-     * @brief static version of read_abacus_orb. A delete-new operation may cause the memory leak, 
-     * it is better to use std::vector to replace the raw pointer.
-     * 
-     * @param ifs [in] ifstream from the orbital file, via `std::ifstream ifs(forb);`
-     * @param elem [out] element symbol
-     * @param ecut [out] planewave energy cutoff
-     * @param nr [out] number of radial grid points
-     * @param dr [out] radial grid spacing
-     * @param nzeta [out] number of zeta functions for each angular momentum
-     * @param radials [out] radial orbitals
-     * @param rank [in] MPI rank
-     */
-    static void read_abacus_orb(std::ifstream& ifs,
-                                std::string& elem,
-                                double& ecut,
-                                int& nr,
-                                double& dr,
-                                std::vector<int>& nzeta,
-                                std::vector<std::vector<double>>& radials,
-                                const int rank = 0);
-
   private:
     double orb_ecut_; //!< energy cutoff as given by the orbital file
 

source/module_basis/module_nao/radial_set.cpp

@@ -7,6 +7,9 @@
 
 #include "module_base/spherical_bessel_transformer.h"
 
+// FIXME: should update with pyabacus
+// #include "module_io/orb_io.h"
+
 RadialSet::~RadialSet()
 {
     delete[] nzeta_;
@@ -252,3 +255,28 @@ void RadialSet::write_abacus_orb(const std::string& file_name, const int rank) c
     }
     file_to.close();
 }
+
+// FIXME: should update with pyabacus
+// void RadialSet::write_abacus_orb(const std::string& forb, const int rank) const
+// {
+//     std::ofstream ofs;
+//     ofs.open(forb, std::ios::out);
+
+//     const double dr = 0.01;
+//     const int nr = static_cast<int>(rcut_max_ / dr) + 1;
+//     std::vector<int> nzeta(lmax_ + 1);
+//     std::copy(nzeta_, nzeta_ + lmax_ + 1, nzeta.begin());
+//     std::vector<std::vector<double>> radials(nchi_, std::vector<double>(nr, 0.0));
+//     int ichi = 0;
+//     for (int l = 0; l <= lmax_; l++)
+//     {
+//         for (int izeta = 0; izeta < nzeta[l]; izeta++)
+//         {
+//             std::copy(chi_[index(l, izeta)].rvalue(), chi_[index(l, izeta)].rvalue() + nr, radials[ichi].begin());
+//             ichi++;
+//         }
+//     }
+
+//     ModuleIO::write_abacus_orb(ofs, symbol_, 100, nr, dr, nzeta, radials, rank);
+//     ofs.close();
+// }

source/module_basis/module_nao/test/CMakeLists.txt

@@ -17,6 +17,7 @@ AddTest(
     ../projgen.cpp
     ../../module_ao/ORB_atomic_lm.cpp
     ../../module_ao/ORB_atomic.cpp
+    ../../../module_io/orb_io.cpp
   LIBS parameter ${math_libs} device base 
 )
 
@@ -29,6 +30,7 @@ AddTest(
     ../numerical_radial.cpp
     ../../module_ao/ORB_atomic_lm.cpp
     ../../module_ao/ORB_atomic.cpp
+    ../../../module_io/orb_io.cpp
   LIBS parameter ${math_libs} device base 
 )
 
@@ -41,6 +43,7 @@ AddTest(
     ../numerical_radial.cpp
     ../../module_ao/ORB_atomic_lm.cpp
     ../../module_ao/ORB_atomic.cpp
+    ../../../module_io/orb_io.cpp
   LIBS parameter ${math_libs} device base 
 )
 
@@ -53,6 +56,7 @@ AddTest(
     ../numerical_radial.cpp
     ../../module_ao/ORB_atomic_lm.cpp
     ../../module_ao/ORB_atomic.cpp
+    ../../../module_io/orb_io.cpp
   LIBS parameter ${math_libs} device base 
 )
 
@@ -65,6 +69,7 @@ AddTest(
     ../numerical_radial.cpp
     ../../module_ao/ORB_atomic_lm.cpp
     ../../module_ao/ORB_atomic.cpp
+    ../../../module_io/orb_io.cpp
  LIBS parameter ${math_libs} device base 
 )
 
@@ -83,6 +88,7 @@ AddTest(
     ../sphbes_radials.cpp
     ../../module_ao/ORB_atomic_lm.cpp
     ../../module_ao/ORB_atomic.cpp
+    ../../../module_io/orb_io.cpp
   LIBS parameter ${math_libs} device base 
 )
 
@@ -103,6 +109,7 @@ AddTest(
     ../two_center_bundle.cpp
     ../two_center_integrator.cpp
     ../real_gaunt_table.cpp
+    ../../../module_io/orb_io.cpp
   LIBS parameter ${math_libs} device base container orb 
 )
 
@@ -132,6 +139,7 @@ AddTest(
     ../radial_set.cpp
     ../numerical_radial.cpp
     ../two_center_bundle.cpp
+    ../../../module_io/orb_io.cpp
   LIBS parameter ${math_libs} device base container orb 
 )
 
@@ -152,6 +160,7 @@ AddTest(
     ../radial_set.cpp
     ../projgen.cpp
     ../numerical_radial.cpp
+    ../../../module_io/orb_io.cpp
   LIBS parameter ${math_libs} device base container orb 
 )
 
@@ -172,6 +181,7 @@ AddTest(
     ../radial_set.cpp
     ../projgen.cpp
     ../numerical_radial.cpp
+    ../../../module_io/orb_io.cpp
   LIBS parameter ${math_libs} device base container orb 
 )
 

source/module_basis/module_nao/test/atomic_radials_test.cpp

@@ -36,10 +36,6 @@ using ModuleBase::SphericalBesselTransformer;
  *
  *  - to_numerical_orbital
  *      - Overwrites the content of a Numerical_Orbital object with the current object.
- * 
- *  - read_abacus_orb
- *     - The static version, reads the radial functions from an abacus file.
- * 
  *                                                                      */
 class AtomicRadialsTest : public ::testing::Test
 {
@@ -268,44 +264,6 @@ TEST_F(AtomicRadialsTest, ToNumericalOrbital)
     }
 }
 
-TEST_F(AtomicRadialsTest, ReadAbacusOrb)
-{
-    std::ifstream ifs(file);
-    std::string elem;
-    double ecut, dr;
-    int nr;
-    std::vector<int> nzeta;
-    std::vector<std::vector<double>> radials;
-    AtomicRadials::read_abacus_orb(ifs, elem, ecut, nr, dr, nzeta, radials);
-    EXPECT_EQ(elem, "Ti");
-    EXPECT_DOUBLE_EQ(ecut, 100.0);
-    EXPECT_EQ(nr, 1001);
-    EXPECT_DOUBLE_EQ(dr, 0.01);
-    EXPECT_EQ(nzeta.size(), 4); // l from 0 to 3
-    EXPECT_EQ(nzeta[0], 4);
-    EXPECT_EQ(nzeta[1], 2);
-    EXPECT_EQ(nzeta[2], 2);
-    EXPECT_EQ(nzeta[3], 1);
-    EXPECT_EQ(radials.size(), 9); // 4 + 2 + 2 + 1
-    for(auto& radial: radials)
-    {
-        EXPECT_EQ(radial.size(), 1001);
-    }
-    EXPECT_EQ(radials[0][0], -1.581711853170e-01);
-    EXPECT_EQ(radials[0][4], -1.583907030513e-01);
-    EXPECT_EQ(radials[0][996], -4.183526380009e-05);
-    EXPECT_EQ(radials[0][1000], 0);
-    EXPECT_EQ(radials[3][0], -1.166292682541e+00);
-    EXPECT_EQ(radials[3][4], -1.164223359672e+00);
-    EXPECT_EQ(radials[3][996], -3.183325576529e-04);
-    EXPECT_EQ(radials[3][1000], 0);
-    EXPECT_EQ(radials[8][0], 0);
-    EXPECT_EQ(radials[8][4], 3.744878535962e-05);
-    EXPECT_EQ(radials[8][996], 7.495357740660e-05);
-    EXPECT_EQ(radials[8][1000], 0);
-    ifs.close();
-}
-
 int main(int argc, char** argv)
 {