Fix DMR spin-channel mismatch in openshell LR spectra

source/source_lcao/module_lr/lr_spectrum.cpp

@@ -100,13 +100,13 @@ ModuleBase::Vector3<std::complex<double>> LR::LR_Spectrum<std::complex<double>>:
         LR_Util::initialize_DMR(DM_trans_real_imag, this->pmat, this->ucell, this->gd_, this->orb_cutoff_);
 
         // real part
-        LR_Util::get_DMR_real_imag_part(DM_trans, DM_trans_real_imag, ucell.nat, 'R');
+        LR_Util::get_DMR_real_imag_part(DM_trans, DM_trans_real_imag, ucell.nat, is, 'R');
         ModuleBase::GlobalFunc::ZEROS(rho_trans_real[0], this->rho_basis.nrxx);
         ModuleGint::cal_gint_rho(DM_trans_real_imag.get_dmr_vec(), 1, rho_trans_real, false);
         // LR_Util::print_grid_nonzero(rho_trans_real[0], this->rho_basis.nrxx, 10, "rho_trans");
 
         // imag part
-        LR_Util::get_DMR_real_imag_part(DM_trans, DM_trans_real_imag, ucell.nat, 'I');
+        LR_Util::get_DMR_real_imag_part(DM_trans, DM_trans_real_imag, ucell.nat, is, 'I');
         ModuleBase::GlobalFunc::ZEROS(rho_trans_imag[0], this->rho_basis.nrxx);
         ModuleGint::cal_gint_rho(DM_trans_real_imag.get_dmr_vec(), 1, rho_trans_imag, false);
         // LR_Util::print_grid_nonzero(rho_trans_imag[0], this->rho_basis.nrxx, 10, "rho_trans");

source/source_lcao/module_lr/utils/lr_util_hcontainer.cpp

@@ -19,6 +19,9 @@ namespace LR_Util
                 {
                     auto ap = dr->find_pair(ia, ja);
                     auto ap_real = dr_real->find_pair(ia, ja);
+                    // under MPI-parallel (2D block-cyclic) HContainer, an atom pair not owned by this rank
+                    // is absent from find_pair() and returns nullptr here; skip it
+                    if (!ap || !ap_real) { continue; }
                     for (int iR = 0;iR < ap->get_R_size();++iR)
                     {
                         // R index may be different between the two HContainers, find by R value instead of R-index
@@ -32,6 +35,39 @@ namespace LR_Util
         }
     }
 
+    void get_DMR_real_imag_part(const module_dm::DensityMatrix<std::complex<double>, std::complex<double>>& DMR,
+        module_dm::DensityMatrix<std::complex<double>, double>& DMR_real,
+        const int& nat,
+        const int& is,
+        const char& type)
+    {
+        assert(is < DMR.get_dmr_vec().size());
+        assert(DMR_real.get_dmr_vec().size() == 1);
+        bool get_imag = (type == 'I' || type == 'i');
+        auto dr = DMR.get_dmr_vec()[is]; //get_dmr_ptr() has bug when is=0
+        auto dr_real = DMR_real.get_dmr_vec()[0];
+        assert(dr != nullptr);
+        assert(dr_real != nullptr);
+        for (int ia = 0;ia < nat;ia++) {
+            for (int ja = 0;ja < nat;ja++)
+            {
+                auto ap = dr->find_pair(ia, ja);
+                auto ap_real = dr_real->find_pair(ia, ja);
+                // under MPI-parallel (2D block-cyclic) HContainer, an atom pair not owned by this rank
+                // is absent from find_pair() and returns nullptr here; skip it
+                if (!ap || !ap_real) { continue; }
+                for (int iR = 0;iR < ap->get_R_size();++iR)
+                {
+                    // R index may be different between the two HContainers, find by R value instead of R-index
+                    auto dR = ap->get_R_index(iR);
+                    auto ptr = ap->get_HR_values(iR).get_pointer();
+                    auto ptr_real = ap_real->get_HR_values(dR.x, dR.y, dR.z).get_pointer();
+                    for (int i = 0;i < ap->get_size();++i) { ptr_real[i] = (get_imag ? ptr[i].imag() : ptr[i].real()); }
+                }
+            }
+        }
+    }
+
     void set_HR_real_imag_part(const hamilt::HContainer<double>& HR_real,
         hamilt::HContainer<std::complex<double>>& HR,
         const int& nat,
@@ -43,6 +79,9 @@ namespace LR_Util
             {
                 auto ap = HR.find_pair(ia, ja);
                 auto ap_real = HR_real.find_pair(ia, ja);
+                // under MPI-parallel (2D block-cyclic) HContainer, an atom pair not owned by this rank
+                // is absent from find_pair() and returns nullptr here; skip it
+                if (!ap || !ap_real) { continue; }
                 for (int iR = 0;iR < ap->get_R_size();++iR)
                 {
                     // R index may be different between the two HContainers, find by R value instead of R-index

source/source_lcao/module_lr/utils/lr_util_hcontainer.h

@@ -42,6 +42,13 @@ namespace LR_Util
         module_dm::DensityMatrix<std::complex<double>, double>& DMR_real,
         const int& nat,
         const char& type = 'R');
+    /// overload: only copy the `is`-th spin channel of DMR (source) into the (single-channel) DMR_real,
+    /// to avoid mixing/overlapping spin channels when DMR has more than one spin channel
+    void get_DMR_real_imag_part(const module_dm::DensityMatrix<std::complex<double>, std::complex<double>>& DMR,
+        module_dm::DensityMatrix<std::complex<double>, double>& DMR_real,
+        const int& nat,
+        const int& is,
+        const char& type = 'R');
     void set_HR_real_imag_part(const hamilt::HContainer<double>& HR_real,
         hamilt::HContainer<std::complex<double>>& HR,
         const int& nat,