Feature : reading and writing of h(r) and dm(r) in npz format (#4066)

* Feature : reading and writing of h(r) and dm(r) in npz format

* fix bug

* udpate makefile

* fix makefile and cmake

---------

Co-authored-by: wenfei-li <liwenfei@gmail.com>

CMakeLists.txt

@@ -37,6 +37,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_CNPY "Enable cnpy usage." OFF)
 
 # enable json support
 if(ENABLE_RAPIDJSON)
@@ -467,6 +468,27 @@ if(ENABLE_DEEPKS)
   add_compile_definitions(__DEEPKS)
 endif()
 
+if (ENABLE_CNPY)
+  find_path(cnpy_SOURCE_DIR
+    cnpy.h
+    HINTS ${libnpy_INCLUDE_DIR}
+  )
+  if(NOT cnpy_SOURCE_DIR)
+    include(FetchContent)
+    FetchContent_Declare(
+      cnpy
+      GIT_REPOSITORY https://github.com/rogersce/cnpy.git
+      GIT_PROGRESS TRUE
+    )
+    FetchContent_MakeAvailable(cnpy)
+  else()
+    include_directories(${cnpy_INCLUDE_DIR})
+  endif()
+  include_directories(${cnpy_SOURCE_DIR})
+  target_link_libraries(${ABACUS_BIN_NAME} cnpy)
+  add_compile_definitions(__USECNPY)
+endif()
+
 function(git_submodule_update)
   if(GIT_SUBMOD_RESULT EQUAL "0")
     message(DEBUG "Submodule init'ed")

docs/advanced/input_files/input-main.md

@@ -1607,6 +1607,20 @@ These variables are used to control the output of properties.
 - **Description**: Whether to print the upper triangular part of the exchange-correlation matrices in **Kohn-Sham orbital representation** (unit: Ry): $\braket{\psi_i|V_\text{xc}^\text{(semi-)local}+V_\text{exx}+V_\text{DFTU}|\psi_j}$ for each k point into files in the directory `OUT.${suffix}`, which is useful for the subsequent GW calculation. (Note that currently DeePKS term is not included. ) The files are named `k-$k-Vxc`, the meaning of `$k`corresponding to k point and spin  is same as [hs_matrix.md](../elec_properties/hs_matrix.md#out_mat_hs).
 - **Default**: False
 
+### out_hr_npz/out_dm_npz
+
+- **Type**: Boolean
+- **Availability**: Numerical atomic orbital basis
+- **Description**: Whether to print Hamiltonian matrices H(R)/density matrics DM(R) in npz format. This feature does not work for gamma-only calculations. Currently only intended for internal usage.
+- **Default**: False
+
+### dm_to_rho
+
+- **Type**: Boolean
+- **Availability**: Numerical atomic orbital basis
+- **Description**: Reads density matrix DM(R) in npz format and creates electron density on grids. This feature does not work for gamma-only calculations. Only supports serial calculations. Currently only intended for internal usage.
+- **Default**: False
+
 ### out_app_flag
 
 - **Type**: Boolean

source/Makefile.Objects

@@ -222,6 +222,7 @@ OBJS_ESOLVER_LCAO=esolver_ks_lcao.o\
       esolver_ks_lcao_elec.o\
       esolver_ks_lcao_tddft.o\
       esolver_ks_lcao_tmpfunc.o\
+      io_npz.o\
 
 OBJS_GINT=gint.o\
       gint_gamma.o\

source/module_base/global_variable.cpp

@@ -283,6 +283,9 @@ bool out_bandgap = false; // QO added for bandgap printing
 int out_interval = 1;    // convert from out_hsR_interval liuyu 2023-04-18
 
 bool out_mat_xc = false; // output Vxc in KS-wfc representation for GW calculation
+bool out_hr_npz = false;
+bool out_dm_npz = false;
+bool dm_to_rho = false; // reads dm in npz format, then prints density in cube format
 
 //==========================================================
 // Deltaspin related

source/module_base/global_variable.h

@@ -313,6 +313,9 @@ extern bool out_bandgap;
 extern int out_interval;
 
 extern bool out_mat_xc; // output Vxc in KS-wfc representation for GW calculation
+extern bool out_hr_npz; //writes h(r) in npz format
+extern bool out_dm_npz; //writes dm(r) in npz format
+extern bool dm_to_rho; //reads in dm(r) and creates density
 
 // Deltaspin related
 extern bool sc_mag_switch; // 0: no deltaspin; 1: constrain atomic magnetic moments;

source/module_elecstate/elecstate_lcao.cpp

@@ -93,6 +93,12 @@ void ElecStateLCAO<std::complex<double>>::psiToRho(const psi::Psi<std::complex<d
     ModuleBase::timer::tick("ElecStateLCAO", "psiToRho");
 
     this->calculate_weights();
+
+// the calculations of dm, and dm -> rho are, technically, two separate functionalities, as we cannot
+// rule out the possibility that we may have a dm from other sources, such as read from file.
+// However, since we are not separating them now, I opt to add a flag to control how dm is obtained as of now
+if(!GlobalV::dm_to_rho)
+{
     this->calEBand();
 
     ModuleBase::GlobalFunc::NOTE("Calculate the density matrix.");
@@ -130,6 +136,7 @@ void ElecStateLCAO<std::complex<double>>::psiToRho(const psi::Psi<std::complex<d
             this->print_psi(psi);
         }
     }
+}
     // old 2D-to-Grid conversion has been replaced by new Gint Refactor 2023/09/25
     //this->loc->cal_dk_k(*this->lowf->gridt, this->wg, (*this->klist));
     for (int is = 0; is < GlobalV::NSPIN; is++)

source/module_esolver/CMakeLists.txt

@@ -18,6 +18,7 @@ if(ENABLE_LCAO)
       esolver_ks_lcao_elec.cpp
       esolver_ks_lcao_tddft.cpp
       esolver_ks_lcao_tmpfunc.cpp
+      io_npz.cpp
   )
 endif()
 

source/module_esolver/esolver_ks.cpp

@@ -390,6 +390,7 @@ void ESolver_KS<T, Device>::run(const int istep, UnitCell& ucell)
 		ModuleBase::timer::tick(this->classname, "run");
 
 		this->before_scf(istep); //Something else to do before the iter loop
+		if(GlobalV::dm_to_rho) return; //nothing further is needed
 
 		ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "INIT SCF");
 
@@ -631,6 +632,21 @@ ModuleIO::Output_Rho ESolver_KS<T, Device>::create_Output_Rho(
 {
 	const int precision = 3;
 	std::string tag = "CHG";
+	if(GlobalV::dm_to_rho)
+	{
+		return ModuleIO::Output_Rho(this->pw_big,
+								this->pw_rhod,
+								is,
+								GlobalV::NSPIN,
+								pelec->charge->rho[is],
+								iter,
+								this->pelec->eferm.get_efval(is),
+								&(GlobalC::ucell),
+								GlobalV::global_out_dir,
+								precision,
+								tag,
+								prefix);
+	}
 	return ModuleIO::Output_Rho(this->pw_big,
 			this->pw_rhod,
 			is,

source/module_esolver/esolver_ks_lcao.cpp

@@ -1116,6 +1116,38 @@ void ESolver_KS_LCAO<TK, TR>::after_scf(const int istep)
     }
 #endif
 
+    if(GlobalV::out_hr_npz)
+    {
+        this->p_hamilt->updateHk(0); // first k point, up spin
+        hamilt::HamiltLCAO<std::complex<double>, double>* p_ham_lcao 
+            = dynamic_cast<hamilt::HamiltLCAO<std::complex<double>, double>*>(this->p_hamilt);
+        std::string zipname = "output_HR0.npz";
+        this->output_mat_npz(zipname,*(p_ham_lcao->getHR()));
+
+        if(GlobalV::NSPIN==2)
+        {
+            this->p_hamilt->updateHk(this->kv.nks/2); // the other half of k points, down spin
+            hamilt::HamiltLCAO<std::complex<double>, double>* p_ham_lcao 
+                = dynamic_cast<hamilt::HamiltLCAO<std::complex<double>, double>*>(this->p_hamilt);
+            zipname = "output_HR1.npz";
+            this->output_mat_npz(zipname,*(p_ham_lcao->getHR()));            
+        }
+    }
+
+    if(GlobalV::out_dm_npz)
+    {
+        const elecstate::DensityMatrix<TK, double>* dm
+            = dynamic_cast<const elecstate::ElecStateLCAO<TK>*>(this->pelec)->get_DM();
+        std::string zipname = "output_DM0.npz";
+        this->output_mat_npz(zipname,*(dm->get_DMR_pointer(1)));
+
+        if(GlobalV::NSPIN==2)
+        {
+            zipname = "output_DM1.npz";
+            this->output_mat_npz(zipname,*(dm->get_DMR_pointer(2)));       
+        }
+    }
+
     if (!md_skip_out(GlobalV::CALCULATION, istep, GlobalV::out_interval))
     {
         this->create_Output_Mat_Sparse(istep).write();

source/module_esolver/esolver_ks_lcao.h

@@ -115,6 +115,9 @@ namespace ModuleESolver
         /// @brief create ModuleIO::Output_Mat_Sparse object to output sparse density matrix of H, S, T, r
         ModuleIO::Output_Mat_Sparse<TK> create_Output_Mat_Sparse(int istep);
 
+        void read_mat_npz(std::string& zipname, hamilt::HContainer<double>& hR);
+        void output_mat_npz(std::string& zipname, const hamilt::HContainer<double>& hR);
+
         /// @brief check if skip the corresponding output in md calculation
         bool md_skip_out(std::string calculation, int istep, int interval);
 

source/module_esolver/esolver_ks_lcao_elec.cpp

@@ -329,6 +329,29 @@ void ESolver_KS_LCAO<TK, TR>::before_scf(int istep)
         ->get_DM()
         ->init_DMR(*(dynamic_cast<hamilt::HamiltLCAO<TK, TR>*>(this->p_hamilt)->getHR()));
 
+    if(GlobalV::dm_to_rho)
+    {
+        std::string zipname = "output_DM0.npz";
+        elecstate::DensityMatrix<TK, double>* dm
+            = dynamic_cast<const elecstate::ElecStateLCAO<TK>*>(this->pelec)->get_DM();
+        this->read_mat_npz(zipname,*(dm->get_DMR_pointer(1)));
+        if(GlobalV::NSPIN == 2)
+        {
+            zipname = "output_DM1.npz";
+            this->read_mat_npz(zipname,*(dm->get_DMR_pointer(2)));
+        }
+
+        this->pelec->psiToRho(*this->psi);
+
+        this->create_Output_Rho(0, istep).write();
+        if(GlobalV::NSPIN == 2)
+        {
+            this->create_Output_Rho(1, istep).write();
+        }
+
+        return;
+    }
+
     // the electron charge density should be symmetrized,
     // here is the initialization
     Symmetry_rho srho;