Fix: init_chg hr support nspin=2

Author: dyzheng

docs/advanced/scf/initialization.md

@@ -11,6 +11,7 @@ In LCAO basis, wavefunction can be read to calculate initial charge density. The
  - `file` : initial charge density from files produced by previous calculations with [`out_chg 1`](../elec_properties/charge.md).
  - `auto`: Abacus first attempts to read the density from a file; if not found, it defaults to using atomic density.
  - `dm` (LCAO only): initial charge density from density matrix files in CSR format. For `nspin=1`, reads `dmrs1_nao.csr`. For `nspin=2` (spin-polarized), reads both `dmrs1_nao.csr` (spin-up) and `dmrs2_nao.csr` (spin-down). These files are generated by previous calculations with [`out_dmr 1`](../elec_properties/density_matrix.md). This method is particularly useful for restarting spin-polarized calculations.
+ - `hr` (LCAO only): initial charge density from Hamiltonian matrix files in CSR format. The Hamiltonian is read from file, then diagonalized to obtain wavefunctions and charge density. For `nspin=1`, reads `hrs1_nao.csr`. For `nspin=2` (spin-polarized), reads both `hrs1_nao.csr` (spin-up) and `hrs2_nao.csr` (spin-down). These files are generated by previous calculations with [`out_mat_hs2 1`](../input_files/input-main.md).
 
 ## Wave function
 `init_wfc` is used for choosing the method of wavefunction coefficient initialization.

source/source_esolver/esolver_ks_lcao.cpp

@@ -182,13 +182,26 @@ void ESolver_KS_LCAO<TK, TR>::before_scf(UnitCell& ucell, const int istep)
         if(PARAM.inp.init_chg == "hr")
         {
             //! 13.1.2) init HR from file
-            std::string hrfile = PARAM.globalv.global_readin_dir + "hrs1_nao.csr";
-            LCAO_domain::init_hr_from_file<TR>(
-                hrfile,
-                dynamic_cast<hamilt::HamiltLCAO<TK, TR>*>(this->p_hamilt)->getHR(),
-                ucell,
-                &(this->pv)
-            );
+            if (PARAM.inp.nspin == 2)
+            {
+                // nspin=2: load spin-up into first half of hRS2, spin-down into second half
+                const std::string hrfile_up = PARAM.globalv.global_readin_dir + "/hrs1_nao.csr";
+                LCAO_domain::init_hr_from_file<TR>(hrfile_up, hamilt_lcao->getHR(), ucell, &(this->pv));
+
+                // switch hR data pointer to spin-down half, then read hrs2
+                auto& hRS2 = hamilt_lcao->getHRS2();
+                hamilt_lcao->getHR()->allocate(hRS2.data() + hRS2.size() / 2, 0);
+                const std::string hrfile_down = PARAM.globalv.global_readin_dir + "/hrs2_nao.csr";
+                LCAO_domain::init_hr_from_file<TR>(hrfile_down, hamilt_lcao->getHR(), ucell, &(this->pv));
+
+                // restore hR to spin-up half (refresh(false) will also do this, but be explicit)
+                hamilt_lcao->getHR()->allocate(hRS2.data(), 0);
+            }
+            else
+            {
+                const std::string hrfile = PARAM.globalv.global_readin_dir + "/hrs1_nao.csr";
+                LCAO_domain::init_hr_from_file<TR>(hrfile, hamilt_lcao->getHR(), ucell, &(this->pv));
+            }
             this->p_hamilt->refresh(false);
             hsolver::HSolverLCAO<TK> hsolver_lcao_obj(&(this->pv), PARAM.inp.ks_solver);
             hsolver_lcao_obj.solve(this->p_hamilt, this->psi[0], this->pelec, *this->dmat.dm,

source/source_lcao/hamilt_lcao.cpp

@@ -537,8 +537,10 @@ void HamiltLCAO<TK, TR>::refresh(bool yes)
         this->refresh_times = 0;
         if (PARAM.inp.nspin == 2)
         {
-            ModuleBase::WARNING_QUIT("HamiltLCAO::refresh",
-                                      "When turning off the refresh flag, the nspin==2 case is not supported yet.");
+            // HR has been loaded from file into both halves of hRS2.
+            // Reset to spin-up; updateHk will switch pointers as needed.
+            this->current_spin = 0;
+            this->hR->allocate(this->hRS2.data(), 0);
         }
     }
 }

source/source_lcao/hamilt_lcao.h

@@ -123,6 +123,9 @@ class HamiltLCAO : public Hamilt<TK>
     }
 #endif
 
+    /// get hRS2 buffer for NSPIN=2 case (spin-up in first half, spin-down in second half)
+    std::vector<TR>& getHRS2() { return this->hRS2; }
+
     /// refresh the status of HR
     void refresh(bool yes) override;
 

source/source_lcao/test/CMakeLists.txt

@@ -5,7 +5,7 @@ remove_definitions(-D__ROCM)
 if(ENABLE_LCAO)
 AddTest(
   TARGET init_dm_from_file_test
-  LIBS parameter ${math_libs} base device cell_info
+  LIBS parameter ${math_libs} base device
   SOURCES test_init_dm_from_file.cpp tmp_mocks.cpp
     ${ABACUS_SOURCE_DIR}/source_estate/module_dm/density_matrix.cpp
     ${ABACUS_SOURCE_DIR}/source_estate/module_dm/density_matrix_io.cpp
@@ -15,6 +15,7 @@ AddTest(
     ${ABACUS_SOURCE_DIR}/source_lcao/module_hcontainer/read_hcontainer.cpp
     ${ABACUS_SOURCE_DIR}/source_lcao/module_hcontainer/func_transfer.cpp
     ${ABACUS_SOURCE_DIR}/source_lcao/module_hcontainer/func_folding.cpp
+    ${ABACUS_SOURCE_DIR}/source_lcao/module_hcontainer/transfer.cpp
     ${ABACUS_SOURCE_DIR}/source_basis/module_ao/parallel_orbitals.cpp
     ${ABACUS_SOURCE_DIR}/source_io/module_output/sparse_matrix.cpp
     ${ABACUS_SOURCE_DIR}/source_io/module_output/csr_reader.cpp

source/source_lcao/test/tmp_mocks.cpp

@@ -1,5 +1,6 @@
 
 #include "source_cell/unitcell.h"
+#include "source_cell/module_neighbor/sltk_grid_driver.h"
 
 // constructor of Atom
 Atom::Atom() {}
@@ -46,3 +47,12 @@ void UnitCell::set_iat2iwt(const int& npol_in)
     }
     return;
 }
+
+// stub for Grid_Driver::Find_atom (used by density_matrix_io.cpp but not exercised in test)
+void Grid_Driver::Find_atom(const UnitCell& ucell,
+                            const ModuleBase::Vector3<double>& tau,
+                            const int& T,
+                            const int& I,
+                            AdjacentAtomInfo* adjs) const
+{
+}