deepmodeling · mohanchen · May 15, 2024 · Apr 30, 2023 · Apr 30, 2023 · May 2, 2023
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -36,6 +36,7 @@ option(COMMIT_INFO "Print commit information in log" ON)
 option(ENABLE_FFT_TWO_CENTER "Enable FFT-based two-center integral method." ON)
 option(ENABLE_GOOGLEBENCH "Enable GOOGLE-benchmark usage." OFF)
 option(ENABLE_RAPIDJSON "Enable rapid-json usage." OFF)
+option(ENABLE_PEXSI "Enable support for PEXSI." OFF)
 
 
 
@@ -208,6 +209,13 @@ if(ENABLE_LCAO)
   if (ENABLE_FFT_TWO_CENTER)
       add_compile_definitions(USE_NEW_TWO_CENTER)
   endif()
+
+  if(ENABLE_PEXSI)
+    find_package(PEXSI REQUIRED)
+    target_link_libraries(${ABACUS_BIN_NAME} ${PEXSI_LIBRARY} ${SuperLU_DIST_LIBRARY} ${ParMETIS_LIBRARY} ${METIS_LIBRARY} pexsi)
+    include_directories(${PEXSI_INCLUDE_DIR} ${ParMETIS_INCLUDE_DIR})
+    add_compile_definitions(__PEXSI)
+  endif()
 else()
   set(ENABLE_DEEPKS OFF)
   set(ENABLE_LIBRI OFF)

diff --git a/cmake/FindPEXSI.cmake b/cmake/FindPEXSI.cmake
@@ -0,0 +1,57 @@
+###############################################################################
+# - Find PEXSI
+# Find PEXSI and its dependencies.
+#
+#  PEXSI_FOUND          - True if pexsi is found.
+#  PEXSI_INCLUDE_DIR    - Where to find pexsi headers.
+#  PEXSI_LIBRARY        - pexsi library.
+#  ParMETIS_INCLUDE_DIR - Where to find pexsi headers.
+#  ParMETIS_LIBRARY     - parmetis library.
+#  METIS_LIBRARY        - metis library.
+#  SuperLU_DIST_LIBRARY - superlu_dist library.
+
+find_path(PEXSI_INCLUDE_DIR
+    NAMES c_pexsi_interface.h
+    HINTS ${PEXSI_DIR}
+    PATH_SUFFIXES "include"
+)
+
+find_library(PEXSI_LIBRARY
+    NAMES pexsi
+    HINTS ${PEXSI_DIR}
+    PATH_SUFFIXES "lib"
+)
+
+find_path(ParMETIS_INCLUDE_DIR
+    NAMES metis.h parmetis.h
+    HINTS ${ParMETIS_DIR}
+    PATH_SUFFIXES "include"
+)
+
+find_library(METIS_LIBRARY
+    NAMES metis
+    HINTS ${ParMETIS_DIR}
+    PATH_SUFFIXES "lib"
+)
+
+find_library(ParMETIS_LIBRARY
+    NAMES parmetis
+    HINTS ${ParMETIS_DIR}
+    PATH_SUFFIXES "lib"
+)
+
+find_library(SuperLU_DIST_LIBRARY
+    NAMES superlu_dist
+    HINTS ${SuperLU_DIST_DIR}
+    PATH_SUFFIXES "lib"
+)
+
+# Handle the QUIET and REQUIRED arguments and
+# set PEXSI_FOUND to TRUE if all variables are non-zero.
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(PEXSI DEFAULT_MSG PEXSI_LIBRARY PEXSI_INCLUDE_DIR ParMETIS_LIBRARY METIS_LIBRARY SuperLU_DIST_LIBRARY)
+
+
+# Copy the results to the output variables and target.
+mark_as_advanced(PEXSI_LIBRARY PEXSI_INCLUDE_DIR ParMETIS_LIBRARY SuperLU_DIST_LIBRARY)
+
diff --git a/docs/advanced/input_files/input-main.md b/docs/advanced/input_files/input-main.md
@@ -382,6 +382,30 @@
     - [qo\_strategy](#qo_strategy)
     - [qo\_screening\_coeff](#qo_screening_coeff)
     - [qo\_thr](#qo_thr)
+  - [PEXSI](#PEXSI)
+    - [pexsi_npole](#pexsi_npole)
+    - [pexsi_inertia](#pexsi_inertia)
+    - [pexsi_nmax](#pexsi_nmax)
+    - [pexsi_comm](#pexsi_comm)
+    - [pexsi_storage](#pexsi_storage)
+    - [pexsi_ordering](#pexsi_ordering)
+    - [pexsi_row_ordering](#pexsi_row_ordering)
+    - [pexsi_nproc](#pexsi_nproc)
+    - [pexsi_symm](#pexsi_symm)
+    - [pexsi_trans](#pexsi_trans)
+    - [pexsi_method](#pexsi_method)
+    - [pexsi_nproc_pole](#pexsi_nproc_pole)
+    - [pexsi_temp](#pexsi_temp)
+    - [pexsi_gap](#pexsi_gap)
+    - [pexsi_delta_e](#pexsi_delta_e)
+    - [pexsi_mu_lower](#pexsi_mu_lower)
+    - [pexsi_mu_upper](#pexsi_mu_upper)
+    - [pexsi_mu](#pexsi_mu)
+    - [pexsi_mu_thr](#pexsi_mu_thr)
+    - [pexsi_mu_expand](#pexsi_mu_expand)
+    - [pexsi_mu_guard](#pexsi_mu_guard)
+    - [pexsi_elec_thr](#pexsi_elec_thr)
+    - [pexsi_zero_thr](#pexsi_zero_thr)
 
 [back to top](#full-list-of-input-keywords)
 
@@ -3549,5 +3573,146 @@ These variables are used to control the usage of QO analysis. Please note presen
 - **Description**: the convergence threshold determining the cutoff of generated orbital. Lower threshold will yield orbital with larger cutoff radius.
 - **Default**: 1.0e-6
 
+## PEXSI
+
+These variables are used to control the usage of PEXSI (Pole Expansion and Selected Inversion) method in calculations.
+
+### pexsi_npole
+
+- **Type**: Integer
+- **Description**: the number of poles used in the pole expansion method, should be a even number.
+- **Default**: 80
+
+### pexsi_inertia
+
+- **Type**: Boolean
+- **Description**: whether inertia counting is used at the very beginning.
+- **Default**: True
+
+### pexsi_nmax
+
+- **Type**: Integer
+- **Description**: maximum number of PEXSI iterations after each inertia counting procedure.
+- **Default**: 80
+
+### pexsi_comm
+
+- **Type**: Boolean
+- **Description**: whether to construct PSelInv communication pattern.
+- **Default**: True
+
+### pexsi_storage
+
+- **Type**: Boolean
+- **Description**: whether to use symmetric storage space used by the Selected Inversion algorithm for symmetric matrices.
+- **Default**: True
+
+### pexsi_ordering
+
+- **Type**: Integer
+- **Description**: ordering strategy for factorization and selected inversion. 0: Parallel ordering using ParMETIS, 1: Sequential ordering using METIS, 2: Multiple minimum degree ordering
+- **Default**: 0
+
+### pexsi_row_ordering
+
+- **Type**: Integer
+- **Description**: row permutation strategy for factorization and selected inversion, 0: No row permutation, 1: Make the diagonal entry of the matrix larger than the off-diagonal entries.
+- **Default**: 1
+
+### pexsi_nproc
+
+- **Type**: Integer
+- **Description**: number of processors for PARMETIS. Only used if pexsi_ordering == 0.
+- **Default**: 1
+
+### pexsi_symm
+
+- **Type**: Boolean
+- **Description**: whether the matrix is symmetric.
+- **Default**: True
+
+### pexsi_trans
+
+- **Type**: Boolean
+- **Description**: whether to factorize the transpose of the matrix.
+- **Default**: False
+
+### pexsi_method
+
+- **Type**: Integer
+- **Description**: the pole expansion method to be used. 1 for Cauchy Contour Integral method, 2 for Moussa optimized method.
+- **Default**: 1
+
+### pexsi_nproc_pole
+
+- **Type**: Integer
+- **Description**: the point parallelizaion of PEXSI. Recommend two points parallelization.
+- **Default**: 1
+
+### pexsi_temp
+
+- **Type**: Real
+- **Description**: temperature in Fermi-Dirac distribution, in Ry, should have the same effect as the smearing sigma when smearing method is set to Fermi-Dirac.
+- **Default**: 0.015
+
+### pexsi_gap
+
+- **Type**: Real
+- **Description**: spectral gap, this can be set to be 0 in most cases.
+- **Default**: 0
+
+### pexsi_delta_e
+
+- **Type**: Real
+- **Description**: an upper bound for the spectral radius of $S^{-1} H$.
+- **Default**: 20
+
+### pexsi_mu_lower
+
+- **Type**: Real
+- **Description**: initial guess of lower bound for mu.
+- **Default**: -10
+
+### pexsi_mu_upper
+
+- **Type**: Real
+- **Description**: initial guess of upper bound for mu.
+- **Default**: 10
+
+### pexsi_mu
+
+- **Type**: Real
+- **Description**: initial guess for mu (for the solver).
+- **Default**: 0
+
+### pexsi_mu_thr
+
+- **Type**: Real
+- **Description**: stopping criterion in terms of the chemical potential for the inertia counting procedure.
+- **Default**: 0.05
+
+### pexsi_mu_expand
+
+- **Type**: Real
+- **Description**: if the chemical potential is not in the initial interval, the interval is expanded by this value.
+- **Default**: 0.3
+
+### pexsi_mu_guard
+
+- **Type**: Real
+- **Description**: safe guard criterion in terms of the chemical potential to reinvoke the inertia counting procedure.
+- **Default**: 0.2
+
+### pexsi_elec_thr
+
+- **Type**: Real
+- **Description**: stopping criterion of the PEXSI iteration in terms of the number of electrons compared to numElectronExact.
+- **Default**: 0.001
+
+### pexsi_zero_thr
+
+- **Type**: Real
+- **Description**: if the absolute value of CCS matrix element is less than this value, it will be considered as zero.
+- **Default**: 1e-10
 
 [back to top](#full-list-of-input-keywords)
diff --git a/source/Makefile b/source/Makefile
@@ -169,6 +169,13 @@ ifdef DeePMD_DIR
     INCLUDES += -I${TensorFlow_INCLUDE_DIR}
 endif
 
+ifdef PEXSI_DIR
+    OBJS_ABACUS += ${OBJS_HSOLVER_PEXSI}
+    INCLUDES += -I${PEXSI_DIR}/include -I${PARMETIS_DIR}/include -I${DSUPERLU_DIR}/include
+    LIBS += -L${PEXSI_DIR}/lib -lpexsi -L${DSUPERLU_DIR}/lib -lsuperlu_dist -L${PARMETIS_DIR}/lib -lparmetis -lmetis
+    HONG += -D__PEXSI
+endif
+
 include Makefile.Objects
 
 #==========================

diff --git a/source/Makefile.Objects b/source/Makefile.Objects
@@ -30,6 +30,7 @@ VPATH=./src_global:\
 ./module_hsolver:\
 ./module_hsolver/kernels:\
 ./module_hsolver/genelpa:\
+./module_hsolver/module_pexsi:\
 ./module_elecstate:\
 ./module_elecstate/kernels:\
 ./module_elecstate/potentials:\
@@ -102,6 +103,7 @@ ${OBJS_VDW}\
 ${OBJS_DFTU}\
 ${OBJS_DELTASPIN}\
 ${OBJS_TENSOR}\
+${OBJS_HSOLVER_PEXSI}\
 
 OBJS_MAIN=main.o\
     driver.o\
@@ -292,7 +294,7 @@ OBJS_HSOLVER=diago_cg.o\
     diago_iter_assist.o\
     math_kernel_op.o\
     dngvd_op.o\
-
+    
 OBJS_HSOLVER_LCAO=hsolver_lcao.o\
       diago_blas.o\
       diago_elpa.o\
@@ -301,6 +303,13 @@ OBJS_HSOLVER_LCAO=hsolver_lcao.o\
       elpa_new_complex.o\
       utils.o\
 
+OBJS_HSOLVER_PEXSI=diago_pexsi.o\
+      pexsi_solver.o\
+      simple_pexsi.o\
+      dist_bcd_matrix.o\
+      dist_ccs_matrix.o\
+      dist_matrix_transformer.o\
+
 OBJS_MD=fire.o\
     langevin.o\
     md_base.o\

diff --git a/source/Makefile.vars b/source/Makefile.vars
@@ -59,6 +59,7 @@ CEREAL_DIR    = /usr/local/include/cereal
 ## To use LIBXC:  set LIBXC_DIR which contains include and lib/libxc.a (>5.1.7)
 ## To use DeePMD: set DeePMD_DIR and TensorFlow_DIR
 ## To use LibRI:  set LIBRI_DIR and LIBCOMM_DIR
+## To use PEXSI: set PEXSI_DIR DSUPERLU_DIR and PARMETIS_DIR
 ##---------------------------------------------------------------------
 
 # LIBTORCH_DIR  = /usr/local
@@ -72,6 +73,10 @@ CEREAL_DIR    = /usr/local/include/cereal
 # LIBRI_DIR     = /public/software/LibRI
 # LIBCOMM_DIR   = /public/software/LibComm
 
+# PEXSI_DIR = /public/software/pexsi
+# DSUPERLU_DIR = /public/software/superlu_dist
+# PARMETIS_DIR    = /public/software/parmetis
+
 ##---------------------------------------------------------------------
 # NP = 14 # It is not supported. use make -j14 or make -j to parallelly compile
 # DEBUG = OFF

diff --git a/source/module_base/global_function.h b/source/module_base/global_function.h
@@ -350,7 +350,7 @@ T ddot_real(
 //==========================================================
 static inline bool IS_COLUMN_MAJOR_KS_SOLVER()
 {
-    return GlobalV::KS_SOLVER=="genelpa" || GlobalV::KS_SOLVER=="scalapack_gvx" || GlobalV::KS_SOLVER=="cusolver" || GlobalV::KS_SOLVER=="cg_in_lcao";
+    return GlobalV::KS_SOLVER=="genelpa" || GlobalV::KS_SOLVER=="scalapack_gvx" || GlobalV::KS_SOLVER=="cusolver" || GlobalV::KS_SOLVER=="cg_in_lcao" || GlobalV::KS_SOLVER=="pexsi";
 }
 
 }//namespace GlobalFunc

diff --git a/source/module_basis/module_ao/ORB_control.cpp b/source/module_basis/module_ao/ORB_control.cpp
@@ -184,7 +184,7 @@ void ORB_control::setup_2d_division(std::ofstream& ofs_running,
     ofs_running << "\n SETUP THE DIVISION OF H/S MATRIX" << std::endl;
 
     // (1) calculate nrow, ncol, nloc.
-    if (ks_solver == "genelpa" || ks_solver == "scalapack_gvx" || ks_solver == "cusolver" || ks_solver == "cg_in_lcao")
+    if (ks_solver == "genelpa" || ks_solver == "scalapack_gvx" || ks_solver == "cusolver" || ks_solver == "cg_in_lcao" || ks_solver == "pexsi")
     {
         ofs_running << " divide the H&S matrix using 2D block algorithms." << std::endl;
 #ifdef __MPI
@@ -205,7 +205,7 @@ void ORB_control::setup_2d_division(std::ofstream& ofs_running,
     bool div_2d;
     if (ks_solver == "lapack" || ks_solver == "cg" || ks_solver == "dav") div_2d = false;
 #ifdef __MPI
-    else if (ks_solver == "genelpa" || ks_solver == "scalapack_gvx" || ks_solver == "cusolver" || ks_solver == "cg_in_lcao") div_2d = true;
+    else if (ks_solver == "genelpa" || ks_solver == "scalapack_gvx" || ks_solver == "cusolver" || ks_solver == "cg_in_lcao" || ks_solver == "pexsi") div_2d = true;
 #endif
     else
     {
@@ -382,7 +382,7 @@ assert(nb2d > 0);
     }
 
     // init blacs context for genelpa
-    if (ks_solver == "genelpa" || ks_solver == "scalapack_gvx" || ks_solver == "cusolver" || ks_solver == "cg_in_lcao")
+    if (ks_solver == "genelpa" || ks_solver == "scalapack_gvx" || ks_solver == "cusolver" || ks_solver == "cg_in_lcao" || ks_solver == "pexsi")
     {
         pv->set_desc(nlocal, nlocal, pv->nrow);
         pv->set_desc_wfc_Eij(nlocal, nbands, pv->nrow);