feat: lr-grad openshell implementation

Author: maki49

source/source_esolver/esolver_lr_lcao_tddft.h

@@ -160,6 +160,9 @@ namespace ModuleESolver
         void init_pot_groundstate(const Charge& chg_gs);
         ct::Tensor solve_zvector_eqation(const int ispin);
         std::vector<ModuleBase::matrix> cal_force(const int ispin);
+        /// open-shell (spin-unrestricted) excited-state force: X holds [up | down] and every
+        /// density matrix has two independent channels
+        std::vector<ModuleBase::matrix> cal_force_openshell();
         void test_force();   // test: reproduce the force of ground state
         elecstate::DensityMatrix<T, double> cal_dm_gs();  ///< ground-state density matrix
 

source/source_lcao/module_lr/Grad/esolver_lr_grad.cpp

@@ -114,7 +114,7 @@ ct::Tensor ModuleESolver::ESolver_LR<T, TR>::solve_zvector_eqation(const int isp
         std::weak_ptr<Exx_LRI<T>>(this->exx_lri), this->exx_info.info_global.hybrid_alpha,
 #endif
         std::weak_ptr<PotHxcLR>(this->pot[ispin]), std::weak_ptr<PotHxcLR>(this->pot_hxc_gs),
-        this->kv, this->paraX_, this->paraC_, this->paraMat_, this->spin_types[ispin]);
+        this->kv, this->paraX_, this->paraC_, this->paraMat_, this->spin_types[ispin], this->openshell);
     ModuleBase::timer::end("ESolver_LR", "solve_zvector_eqation");
     return Z;
 }
@@ -123,6 +123,7 @@ template<typename T, typename TR>
 std::vector<ModuleBase::matrix> ModuleESolver::ESolver_LR<T, TR>::cal_force(const int ispin)
 {
     if (PARAM.inp.test_force && ispin == 0) { this->test_force(); }
+    if (this->openshell) { return this->cal_force_openshell(); }
 
     const ct::Tensor& Z = this->solve_zvector_eqation(ispin);
 
@@ -302,7 +303,10 @@ std::vector<ModuleBase::matrix> ModuleESolver::ESolver_LR<T, TR>::cal_force(cons
             const auto& Ds_gs = LR_Util::get_exx_Ds_spin1(dm_gs, (*this->ucell_), this->kv, this->paraMat_);    // returns 0.5*D[0]
             const auto& Ds_relaxed_diff = LR_Util::get_exx_Ds_spin1(relaxed_diff_dm, (*this->ucell_), this->kv, this->paraMat_);   // returns 0.5*D[0]
             // LR_Util::print_CV(Ds_relaxed_diff, "Ds_relaxed_diff for EXX force");
-            ModuleBase::matrix force_exx_gs_relaxed_diff = lr_force.cal_force_exx_gs_dm_relaxed_diff(Ds_gs, Ds_relaxed_diff, alpha * 4.0);  // cancel the two 0.5s in Ds
+            // `get_exx_Ds_spin1` feeds `split_m2D_ktoR(..., nspin=1)`, which reads only channel 0
+            // with a 0.5 prefactor. For `dm_gs` that channel is $D^\text{gs}_\uparrow$ at nspin=2
+            // but the spin-summed $D^\text{gs}$ at nspin=1, i.e. twice as large.
+            ModuleBase::matrix force_exx_gs_relaxed_diff = lr_force.cal_force_exx_gs_dm_relaxed_diff(Ds_gs, Ds_relaxed_diff, alpha * 4.0) * gs_dm_channel_factor();  // cancel the two 0.5s in Ds
             if (PARAM.inp.test_force)
                 ModuleIO::print_force(GlobalV::ofs_running, (*this->ucell_), "EXX GS-(T+Z) FORCE (eV/Angstrom)", force_exx_gs_relaxed_diff, false);
             force_hamiltgs_relaxed_diff += force_exx_gs_relaxed_diff;
@@ -312,7 +316,7 @@ std::vector<ModuleBase::matrix> ModuleESolver::ESolver_LR<T, TR>::cal_force(cons
                 // test H[T] force (Z=0), EXX part
                 const auto& Ds_diff = LR_Util::get_exx_Ds_spin1(diff_dm, (*this->ucell_), this->kv, this->paraMat_);   // returns 0.5*D[0]
                 GlobalV::ofs_running << "========== [TEST H_GS-(T) force (Z=0), EXX part] ===========" << std::endl;
-                ModuleBase::matrix force_exx_gs_diff = lr_force.cal_force_exx_gs_dm_relaxed_diff(Ds_gs, Ds_diff, alpha * 4.0);  // cancel the two 0.5s in Ds
+                ModuleBase::matrix force_exx_gs_diff = lr_force.cal_force_exx_gs_dm_relaxed_diff(Ds_gs, Ds_diff, alpha * 4.0) * gs_dm_channel_factor();  // cancel the two 0.5s in Ds
                 ModuleIO::print_force(GlobalV::ofs_running, (*this->ucell_), "H_GS-T EXX FORCE (Z=0) (eV/Angstrom)", force_exx_gs_diff, false);
                 GlobalV::ofs_running << "========== [\\TEST H_GS-(T) force (Z=0), EXX part] ===========" << std::endl;
             }
@@ -327,6 +331,171 @@ std::vector<ModuleBase::matrix> ModuleESolver::ESolver_LR<T, TR>::cal_force(cons
     return forces;
 }
 
+template<typename T, typename TR>
+std::vector<ModuleBase::matrix> ModuleESolver::ESolver_LR<T, TR>::cal_force_openshell()
+{
+    ModuleBase::TITLE("ESolver_LR", "cal_force_openshell");
+    ModuleBase::timer::start("ESolver_LR", "cal_force");
+
+    // Open shell: there is a single eigenproblem whose vector is the concatenation
+    // [up-block | down-block], and every density matrix has two independent channels.
+    // The spin-orbital formulas apply verbatim -- unlike the closed-shell singlet/triplet
+    // algorithm, X here is normalized over BOTH channels, so it carries no implicit sqrt(2)
+    // and none of the collapsed 2/4 factors are needed.
+    const ct::Tensor& Z = this->solve_zvector_eqation(0);
+
+    const std::vector<int> ld_x = { this->nk * this->paraX_[0].get_local_size(),
+                                    this->nk * this->paraX_[1].get_local_size() };
+    const std::vector<int> off_x = { 0, ld_x[0] };
+    std::vector<psi::Psi<T>> c_spin;
+    for (int is : {0, 1}) { c_spin.push_back(LR_Util::get_psi_spin(*this->psi_ks, is, this->nk)); }
+
+    LR_Force<T> lr_force((*this->ucell_), this->kv.kvec_d, this->paraMat_,
+        *this->pw_rhod, *this->pw_rho, this->locpp, this->sf, this->gd, this->two_center_bundle_
+#ifdef __EXX
+        , std::weak_ptr<Exx_LRI<T>>(this->exx_lri), this->exx_info.info_global.hybrid_alpha
+#endif
+    );
+    GlobalV::ofs_running << "Start to calculate excited-state force of updown (open shell)" << std::endl;
+
+    std::vector<ModuleBase::matrix> forces(this->nstates);
+    for (int istate = 0;istate < this->nstates;++istate)
+    {
+        const int offset = istate * this->nloc_per_state;
+        const T* const X_istate = this->X[0].template data<T>() + offset;
+        const T* const Z_istate = Z.template data<T>() + offset;
+
+        // 1. the k-space blocks of each spin channel
+        std::vector<std::vector<ct::Tensor>> dmx_k(2), dmdiff_k(2), relaxed_k(2);
+        for (int is : {0, 1})
+        {
+            dmx_k[is] = cal_dm_trans_pblas(X_istate + off_x[is], this->paraX_[is], c_spin[is], this->paraC_,
+                this->nbasis, this->nocc[is], this->nvirt[is], this->paraMat_);
+            dmdiff_k[is] = cal_dm_diff_pblas(X_istate + off_x[is], this->paraX_[is], c_spin[is], this->paraC_,
+                this->nbasis, this->nocc[is], this->nvirt[is], this->paraMat_);
+            std::vector<ct::Tensor> dmz_k = cal_dm_trans_pblas(Z_istate + off_x[is], this->paraX_[is], c_spin[is],
+                this->paraC_, this->nbasis, this->nocc[is], this->nvirt[is], this->paraMat_);
+            for (auto& d : dmz_k) { LR_Util::matsym(d.template data<T>(), this->nbasis, this->paraMat_); }
+            relaxed_k[is] = dmdiff_k[is] + dmz_k;
+        }
+
+        // 2. $D^X$. Complex and UN-symmetrized first (the EXX kernel needs the full
+        //    non-symmetric $D^X$), then the real symmetrized copy for the grid Hxc force --
+        //    `build_dm_from_dmk_spin` symmetrizes IN PLACE, hence the ordering.
+        auto dm_trans = LR_Util::build_dm_from_dmk_spin<T, T>(dmx_k,
+            this->paraMat_, this->nk, this->kv.kvec_d, (*this->ucell_), this->gd, this->orb_cutoff_);
+        LR_Util::transpose_DMR(dm_trans, (*this->ucell_).nat);
+        auto dm_trans_real = LR_Util::build_dm_from_dmk_spin<T, double>(dmx_k,
+            this->paraMat_, this->nk, this->kv.kvec_d, (*this->ucell_), this->gd, this->orb_cutoff_,
+            /*symmetrize=*/true);
+        LR_Util::transpose_DMR(dm_trans_real, (*this->ucell_).nat);
+
+        // 3. the relaxed difference density matrix $T+D^Z$
+        const elecstate::DensityMatrix<T, T>& relaxed_diff_dm =
+            LR_Util::build_dm_from_dmk_spin<T, T>(relaxed_k,
+                this->paraMat_, this->nk, this->kv.kvec_d, (*this->ucell_), this->gd, this->orb_cutoff_);
+        elecstate::DensityMatrix<T, double> relaxed_diff_dm_real(&this->paraMat_, 2, this->kv.kvec_d, this->nk);
+        LR_Util::initialize_DMR(relaxed_diff_dm_real, this->paraMat_, (*this->ucell_), this->gd, this->orb_cutoff_);
+        LR_Util::get_DMR_real_imag_part(relaxed_diff_dm, relaxed_diff_dm_real, 'R');
+
+        // 4. the energy-weighted density matrix
+        std::weak_ptr<PotHxcLR> pot_weak = this->pot[0];
+        std::weak_ptr<PotHxcLR> pot_hxc_gs_weak = this->pot_hxc_gs;
+#ifdef __EXX
+        std::weak_ptr<Exx_LRI<T>> exx_lri_weak = this->exx_lri;
+#endif
+        const std::vector<std::vector<ct::Tensor>>& edm_k =
+            cal_edm_from_XZ_istate_openshell(X_istate, Z_istate,
+                this->pelec->ekb.c[istate], this->eig_ks.c, dm_trans,
+                *this->psi_ks, this->nspin, this->nbasis, this->nocc, this->nvirt,
+                (*this->ucell_), this->orb_cutoff_,
+#ifdef __EXX
+                exx_lri_weak, this->exx_info.info_global.hybrid_alpha,
+#endif
+                pot_weak, pot_hxc_gs_weak,
+                this->kv, this->gd, this->paraX_, this->paraC_, this->paraMat_, this->xc_kernel);
+        elecstate::DensityMatrix<T, double> edm_real = LR_Util::build_dm_from_dmk_spin<T, double>(edm_k,
+            this->paraMat_, this->nk, this->kv.kvec_d, (*this->ucell_), this->gd, this->orb_cutoff_,
+            /*symmetrize=*/true);
+
+        // 5. the force terms
+        ModuleBase::matrix force_hxc_dmtrans = lr_force.cal_force_hxc_dmtrans(dm_trans_real, *this->pot[0]);
+        if (PARAM.inp.test_force)
+            ModuleIO::print_force(GlobalV::ofs_running, (*this->ucell_), "HXC DMTRANS FORCE (eV/Angstrom)", force_hxc_dmtrans, false);
+
+        const elecstate::DensityMatrix<T, double>& dm_gs = this->cal_dm_gs();
+
+        // the $g^{xc}$ half of $\partial_x K[D^X]D^X$, i.e. the derivative of the xc kernel
+        // through the ground-state density. Only for local kernels.
+        if (LR_Util::has_local_xc(this->xc_kernel))
+        {
+            PotGradXCLR pot_grad(this->pot_hxc_gs->xc_kernel_components(), this->pot_hxc_gs->get_rho_basis(),
+                (*this->ucell_), this->pot_hxc_gs->nrxx, /*triplet=*/false);
+            ModuleBase::matrix force_gxc_dmtrans =
+                lr_force.cal_force_gxc_dmtrans_openshell(dm_trans_real, dm_gs, pot_grad);
+            if (PARAM.inp.test_force)
+                ModuleIO::print_force(GlobalV::ofs_running, (*this->ucell_), "GXC DMTRANS FORCE (eV/Angstrom)", force_gxc_dmtrans, false);
+            force_hxc_dmtrans += force_gxc_dmtrans;
+        }
+
+        ModuleBase::matrix force_hamiltgs_relaxed_diff = lr_force.cal_force_hamilt_gs_dm_relaxed_diff(
+            relaxed_diff_dm_real, dm_gs, false, this->pot_hxc_gs.get());
+        if (PARAM.inp.test_force)
+            ModuleIO::print_force(GlobalV::ofs_running, (*this->ucell_), "H_GS-(T+Z) FORCE (without EXX) (eV/Angstrom)", force_hamiltgs_relaxed_diff, false);
+
+        ModuleBase::matrix force_overlap_edm = lr_force.cal_force_overlap_edm(edm_real);
+        if (PARAM.inp.test_force)
+            ModuleIO::print_force(GlobalV::ofs_running, (*this->ucell_), "OVERLAP-EDM FORCE (eV/Angstrom)", force_overlap_edm, false);
+
+#ifdef __EXX
+        const double& alpha = this->exx_info.info_global.hybrid_alpha;
+        // Exchange is spin-diagonal, so each channel is done independently and summed.
+        // `get_exx_Ds_gs` returns the channels unscaled (SPIN_multiple = 1 at nspin=2), unlike
+        // the closed-shell `get_exx_Ds_spin1` which returns 0.5*D. Counting the closed-shell
+        // `alpha*4.0` back down for both terms:
+        //   $D^XD^X$: each slot goes 0.5*D^X_tot -> D^X_is, i.e. x2 each, and the explicit
+        //     sum over is adds another x2 -- but $D^X_\text{tot}=\sqrt2 D^X_\sigma$ eats one,
+        //     so 4/(2*2) * ... = `alpha`.
+        //   $D^\text{gs}(T{+}D^Z)$: the left slot goes 0.5*D_up -> D_is (x2); the right slot
+        //     goes 0.5*(T+Z)_tot = (T+Z)_up -> (T+Z)_is (x1, no sqrt2 here); the explicit sum
+        //     over is adds x2. So 4/(2*1*2) = `alpha` as well -- NOT `2*alpha`: the earlier
+        //     comment forgot that the closed-shell right slot is already the spin SUM, which is
+        //     exactly what the `for (is)` loop below now supplies.
+        if (LR::exx_kernel_list().count(this->xc_kernel))
+        {
+            const auto& Ds_trans = LR_Util::get_exx_Ds_gs(dm_trans, (*this->ucell_), this->kv, this->paraMat_);
+            ModuleBase::matrix force_exx_dmtrans(this->ucell_->nat, 3);
+            for (int is : {0, 1})
+            {
+                force_exx_dmtrans += lr_force.cal_force_exx_dm_trans(Ds_trans[is], alpha, std::to_string(is));
+            }
+            if (PARAM.inp.test_force)
+                ModuleIO::print_force(GlobalV::ofs_running, (*this->ucell_), "EXX DMTRANS FORCE (eV/Angstrom)", force_exx_dmtrans, false);
+            force_hxc_dmtrans += force_exx_dmtrans;
+        }
+        if (LR::exx_kernel_list().count(PARAM.inp.dft_functional))
+        {
+            const auto& Ds_gs = LR_Util::get_exx_Ds_gs(dm_gs, (*this->ucell_), this->kv, this->paraMat_);
+            const auto& Ds_relaxed_diff = LR_Util::get_exx_Ds_gs(relaxed_diff_dm, (*this->ucell_), this->kv, this->paraMat_);
+            ModuleBase::matrix force_exx_gs_relaxed_diff(this->ucell_->nat, 3);
+            for (int is : {0, 1})
+            {
+                force_exx_gs_relaxed_diff += lr_force.cal_force_exx_gs_dm_relaxed_diff(
+                    Ds_gs[is], Ds_relaxed_diff[is], alpha, std::to_string(is));
+            }
+            if (PARAM.inp.test_force)
+                ModuleIO::print_force(GlobalV::ofs_running, (*this->ucell_), "EXX GS-(T+Z) FORCE (eV/Angstrom)", force_exx_gs_relaxed_diff, false);
+            force_hamiltgs_relaxed_diff += force_exx_gs_relaxed_diff;
+        }
+#endif
+        forces[istate] = force_hxc_dmtrans + force_hamiltgs_relaxed_diff + force_overlap_edm;
+    }
+    ModuleBase::timer::end("ESolver_LR", "cal_force");
+    print_force(forces, std::cout);
+    print_force(forces, GlobalV::ofs_running);
+    return forces;
+}
+
 template<typename T, typename TR>
 elecstate::DensityMatrix<T, double> ModuleESolver::ESolver_LR<T, TR>::cal_dm_gs()
 {

source/source_lcao/module_lr/Grad/force/lr_force.cpp

@@ -7,8 +7,13 @@
 // #include "source_lcao/module_lr/utils/lr_util_hcontainer.h"
 namespace LR
 {
-    /// `dm_gs` carries the ground-state occupations, at nspin=1 they are 2 for fully occupied bands.
-    inline double gs_dm_channel_factor() { return (PARAM.inp.nspin == 1) ? 0.5 : 1.0; }
+    /// The LR density matrices ($D^X$, $T+D^Z$, EDM) carry one channel in the closed-shell
+    /// singlet/triplet algorithm and two independent channels in the open-shell one.
+    template<typename TK>
+    inline bool is_openshell_dm(const elecstate::DensityMatrix<TK, double>& dm)
+    {
+        return dm.get_DMR_vector().size() == 2;
+    }
 
     template<typename TK>
     Charge LR_Force<TK>::dm_to_charge(const elecstate::DensityMatrix<TK, double>& dm)
@@ -58,6 +63,10 @@ namespace LR
         const PotHxcLR* pot_hxc_gs)
     {
         const bool with_ewald = reproduce_gs;
+        // Two independent spin channels in `relax_diff_dm` <=> open-shell (spin-unrestricted) LR.
+        // The closed-shell singlet/triplet algorithm always builds a single-channel LR density
+        // matrix, even at nspin=2.
+        const bool openshell = is_openshell_dm(relax_diff_dm);
         const Charge chr_diff_relaxed = dm_to_charge(relax_diff_dm);
 
         // 1. local pp (Hellmann-Feynman)(fvl_dvl) + ewald + core correction (+ self-consistent charge)
@@ -122,6 +131,27 @@ namespace LR
             //`cal_pulay_fs` calculates only one spin channel because `relax_diff_dm` has only one.
             PulayForceStress::cal_pulay_fs(1/*nspin*/, fhxc_dvhxc, stress_tmp,
                 dm_gs, this->ucell_, &pot_hxc_relaxed_diff, true, false);
+            fhxc_dvhxc *= gs_dm_channel_factor();
+        }
+        else if (openshell)
+        {
+            // $v_\sigma=\sum_{\sigma'}f^{\sigma\sigma'}\rho^{T+Z}_{\sigma'}$ (+ the full Hartree),
+            // contracted with $D^\text{gs}_\sigma$. No factor 2 and no `gs_dm_channel_factor`:
+            // the two ground-state channels are summed explicitly by `cal_gint_fvl` below, and
+            // each carries occupation 1 rather than 2.
+            constexpr int nspin_dm = 2;
+            std::vector<ModuleBase::matrix> v_lin(nspin_dm, ModuleBase::matrix(1, this->rhopw_.nrxx));
+            for (int sl = 0; sl < nspin_dm; ++sl)
+            {
+                for (int sr = 0; sr < nspin_dm; ++sr)
+                {
+                    double* rho_in[1] = { const_cast<double*>(chr_diff_relaxed.rho[sr]) };
+                    pot_hxc_gs->cal_v_eff(rho_in, this->ucell_, v_lin[sl], { sl, sr });
+                }
+            }
+            std::vector<const double*> vr_eff(nspin_dm);
+            for (int is = 0; is < nspin_dm; ++is) { vr_eff[is] = v_lin[is].c; }
+            ModuleGint::cal_gint_fvl(nspin_dm, vr_eff, dm_gs.get_DMR_vector(), true, false, &fhxc_dvhxc, &stress_tmp);
         }
         else
         {
@@ -130,9 +160,9 @@ namespace LR
             pot_hxc_gs->cal_v_eff(rho_in, this->ucell_, v_lin);
             std::vector<const double*> vr_eff = { v_lin.c };
             ModuleGint::cal_gint_fvl(1, vr_eff, dm_gs.get_DMR_vector(), true, false, &fhxc_dvhxc, &stress_tmp);
+            fhxc_dvhxc *= 2;   // for the two channels of the ground-state dm.
+            fhxc_dvhxc *= gs_dm_channel_factor();
         }
-        if(!reproduce_gs) {fhxc_dvhxc *= 2;} // for the two channels of the ground-state dm.
-        fhxc_dvhxc *= gs_dm_channel_factor();
 
         // 4. kinetic (Pulay)
         std::vector<hamilt::HContainer<double>> dT = cal_hs_grad('T', this->ucell_, this->pv_, this->gd_, this->two_center_bundle_);
@@ -165,6 +195,13 @@ namespace LR
         // Verified on H2/SZ against the analytic 4-center derivative: 4*sum D^sym P = 16.8653548
         // vs 2*d(ai|ia)/dz = 16.865355 eV/Ang (7 digits).
         const double pulay_to_total_sym = 2.0;
+        // Open shell: the kernel couples the channels, so the density has to be built per channel
+        // and the potential accumulated over the summed spin before contracting with $D^X_\sigma$.
+        // `pulay_to_total_sym` is a Pulay -> Pulay+Hellmann-Feynman factor and is spin-independent.
+        if (is_openshell_dm(dm_trans))
+        {
+            return PulayForceStress::cal_pulay_fs_openshell(dm_trans, this->ucell_, &pot_hxc) * pulay_to_total_sym;
+        }
         return PulayForceStress::cal_pulay_fs(dm_trans, this->ucell_, &pot_hxc) * pulay_to_total_sym;
     }
 
@@ -201,6 +238,31 @@ namespace LR
         return f;
     }
 
+    template<typename TK>
+    ModuleBase::matrix LR_Force<TK>::cal_force_gxc_dmtrans_openshell(
+        const elecstate::DensityMatrix<TK, double>& dm_trans,
+        const elecstate::DensityMatrix<TK, double>& dm_gs, const PotGradXCLR& pot_grad)
+    {
+        // Same term as `cal_force_gxc_dmtrans`, spin-resolved: the free index $\tau$ (spin channel)
+        // of $v^{(2)}_\tau$ is contracted with $D^\text{gs}_\tau$, and `cal_gint_fvl` does the
+        // $\sum_\tau$. No `gs_dm_channel_factor` here -- both ground-state channels are summed
+        // explicitly and each carries occupation 1.
+        constexpr int nspin_dm = 2;
+        assert(dm_trans.get_DMR_vector().size() == nspin_dm);
+        const Charge chr_x = dm_to_charge(dm_trans);
+        const double* rho_in[nspin_dm] = { chr_x.rho[0], chr_x.rho[1] };
+
+        std::vector<ModuleBase::matrix> v2(nspin_dm, ModuleBase::matrix(1, this->rhopw_.nrxx));
+        for (int tau = 0; tau < nspin_dm; ++tau) { pot_grad.cal_v_eff_openshell(rho_in, this->ucell_, v2[tau], tau); }
+
+        ModuleBase::matrix f(this->ucell_.nat, 3);
+        ModuleBase::matrix stress_tmp;
+        std::vector<const double*> vr_eff(nspin_dm);
+        for (int is = 0; is < nspin_dm; ++is) { vr_eff[is] = v2[is].c; }
+        ModuleGint::cal_gint_fvl(nspin_dm, vr_eff, dm_gs.get_DMR_vector(), true, false, &f, &stress_tmp);
+        return f;
+    }
+
 #ifdef __EXX
     template<typename TK>
     ModuleBase::matrix LR_Force<TK>::cal_force_exx_dm_trans(