Add KSStateToVariable method for handling occupation numbers and wavefunctions in RDMFT_LCAO

Author: Kai Luo

source/source_esolver/directmin_problems/rdmft_lcao.cpp

@@ -305,6 +305,106 @@ void RDMFT_LCAO<TK, TR>::VariableToOccNumWfc(const Variable& x,
     }
 }
 
+template <typename TK, typename TR>
+void RDMFT_LCAO<TK, TR>::KSStateToVariable(const ModuleBase::matrix& occ_num,
+                                           const psi::Psi<TK>& wfc,
+                                           Variable& x) const
+{
+    const int required_components = this->nk_total;
+    const int expected_elements = required_components * 2;
+    if (x.Getnumofelements() != expected_elements)
+    {
+        throw std::runtime_error("KSStateToVariable received a variable with unexpected component count");
+    }
+
+    if (occ_num.nr != this->nk_total || occ_num.nc != this->nbands_total)
+    {
+        throw std::runtime_error("KSStateToVariable received occupation matrix with inconsistent shape");
+    }
+
+    const int occ_offset = required_components;
+    for (int ik = 0; ik < this->nk_total; ++ik)
+    {
+        Element& occ_elem = x.GetElement(occ_offset + ik);
+        if (occ_elem.Getrow() != this->nbands_total || occ_elem.Getcol() != 1 || occ_elem.Getiscomplex())
+        {
+            throw std::runtime_error("KSStateToVariable encountered an occupation component with unexpected shape");
+        }
+
+        double* occ_buffer = occ_elem.ObtainWriteEntireData();
+        for (int ib = 0; ib < this->nbands_total; ++ib)
+        {
+            occ_buffer[ib] = occ_num(ik, ib);
+        }
+    }
+
+    const int local_basis_count = this->pv.get_row_size();
+    const int local_band_slots = this->pv.get_col_size();
+    const int local_band_count = this->pv.ncol_bands;
+    const int band_limit = std::min(local_band_count, local_band_slots);
+
+    if (wfc.get_nk() != this->nk_total || wfc.get_nbands() != local_band_count || wfc.get_nbasis() != local_basis_count)
+    {
+        throw std::runtime_error("KSStateToVariable received wavefunction object with incompatible dimensions");
+    }
+
+    for (int ik = 0; ik < required_components; ++ik)
+    {
+        Element& wf_elem = x.GetElement(ik);
+        if (wf_elem.Getrow() != this->nbasis_total || wf_elem.Getcol() != this->nbands_total)
+        {
+            throw std::runtime_error("KSStateToVariable encountered a wavefunction block with unexpected global shape");
+        }
+
+        double* raw_data = wf_elem.ObtainWriteEntireData();
+        const bool is_complex = wf_elem.Getiscomplex();
+        const int element_length = wf_elem.Getlength();
+        std::fill(raw_data, raw_data + element_length, 0.0);
+        std::complex<double>* complex_buffer = reinterpret_cast<std::complex<double>*>(raw_data);
+
+    wfc.fix_k(ik);
+
+        for (int ib_local = 0; ib_local < band_limit; ++ib_local)
+        {
+            const int band_global = this->pv.local2global_col(ib_local);
+            if (band_global < 0 || band_global >= this->nbands_total)
+            {
+                continue;
+            }
+
+            wfc.fix_b(ib_local);
+
+            for (int ir_local = 0; ir_local < local_basis_count; ++ir_local)
+            {
+                const int basis_global = this->pv.local2global_row(ir_local);
+                if (basis_global < 0 || basis_global >= this->nbasis_total)
+                {
+                    continue;
+                }
+
+                const int linear_index = basis_global + this->nbasis_total * band_global;
+                const TK value = wfc(ir_local);
+
+                if (is_complex)
+                {
+                    complex_buffer[linear_index] = std::complex<double>(std::real(value), std::imag(value));
+                }
+                else
+                {
+                    raw_data[linear_index] = static_cast<double>(std::real(value));
+                }
+            }
+        }
+    }
+
+    if (this->ucell_ref_ == nullptr)
+    {
+        throw std::runtime_error("KSStateToVariable requires a valid UnitCell reference");
+    }
+
+    this->rdmft_solver.update_elec(*this->ucell_ref_, occ_num, wfc);
+}
+
 
 
 

source/source_esolver/directmin_problems/rdmft_lcao.h

@@ -88,6 +88,10 @@ class RDMFT_LCAO : public Problem, public ESolver_KS_LCAO<TK, TR>
         ModuleBase::matrix &occ_num, 
         psi::Psi<TK> &wfc) const;
 
+    void KSStateToVariable(const ModuleBase::matrix& occ_num,
+                           const psi::Psi<TK>& wfc,
+                           Variable& x) const;
+
 
 };
 

source/source_esolver/esolver_directmin_lcao.cpp

@@ -15,6 +15,8 @@
 #include "Manifolds/Euclidean.h"
 #include "Manifolds/CStiefel.h"
 
+#include <string>
+
 // using namespace ROPTLITE;
 
 // #include "source_lcao/module_rdmft/rdmft.h"
@@ -46,6 +48,25 @@
 #include "source_lcao/module_ri/exx_opt_orb.h"
 #endif
 
+namespace
+{
+template <typename TK, typename TR>
+bool seed_variable_from_ks_solver(ModuleESolver::ESolver_KS_LCAO<TK, TR>* ks_solver,
+                                  ModuleESolver::RDMFT_LCAO<TK, TR>* rdmft_prob,
+                                  ROPTLITE::Variable& variable)
+{
+    if (ks_solver == nullptr || rdmft_prob == nullptr)
+    {
+        return false;
+    }
+
+    rdmft_prob->KSStateToVariable(ks_solver->get_elecstate().wg,
+                                  ks_solver->get_wavefunctions(),
+                                  variable);
+    return true;
+}
+} // namespace
+
 namespace ModuleESolver
 {
 
@@ -361,14 +382,73 @@ void ESolver_DirectMin_LCAO::setup_solver(UnitCell& ucell, const Input_para& inp
     // for the moment, use random use random orbitals and fixed weights
     else if (inp.directmin_objective == "rdmft") 
     {
-        X = this->prob->GetDomain()->RandominManifold();
-    // check if directmin_init_method is "ks", then we need to get the KS wavefunctions and occupations
-    // (was rdmft_init_method)
-        // from the previous SCF calculation
+        Manifold* domain = this->prob->GetDomain();
+        if (domain == nullptr)
+        {
+            ModuleBase::WARNING_QUIT("ESolver_DirectMin_LCAO", "RDMFT problem domain is undefined");
+        }
+
+        Variable initial_variable = domain->RandominManifold();
+        bool seeded = false;
+        std::string seed_error;
+
         if (this->init_method_ == "ks")
         {
-            
+            ModuleESolver::ESolver* p_esolver = nullptr;
+            if (PARAM.globalv.gamma_only_local)
+            {
+                p_esolver = new ESolver_KS_LCAO<double, double>();
+            }
+            else if (PARAM.inp.nspin < 4)
+            {
+                p_esolver = new ESolver_KS_LCAO<std::complex<double>, double>();
+            }
+            else
+            {
+                p_esolver = new ESolver_KS_LCAO<std::complex<double>, std::complex<double>>();
+            }
+
+            try
+            {
+                p_esolver->before_all_runners(ucell, inp);
+                p_esolver->runner(ucell, 0);
+
+                seeded = seed_variable_from_ks_solver(dynamic_cast<ESolver_KS_LCAO<double, double>*>(p_esolver),
+                                                       dynamic_cast<RDMFT_LCAO<double, double>*>(this->prob),
+                                                       initial_variable)
+                    || seed_variable_from_ks_solver(dynamic_cast<ESolver_KS_LCAO<std::complex<double>, double>*>(p_esolver),
+                                                    dynamic_cast<RDMFT_LCAO<std::complex<double>, double>*>(this->prob),
+                                                    initial_variable)
+                    || seed_variable_from_ks_solver(dynamic_cast<ESolver_KS_LCAO<std::complex<double>, std::complex<double>>*>(p_esolver),
+                                                    dynamic_cast<RDMFT_LCAO<std::complex<double>, std::complex<double>>*>(this->prob),
+                                                    initial_variable);
+            }
+            catch (const std::exception& ex)
+            {
+                seed_error = ex.what();
+            }
+
+            ModuleESolver::clean_esolver(p_esolver, false);
+
+            if (!seeded)
+            {
+                if (!seed_error.empty())
+                {
+                    const std::string message = "Unable to seed DirectMin variable from KS solver: " + seed_error;
+                    ModuleBase::WARNING("ESolver_DirectMin_LCAO", message.c_str());
+                }
+                else
+                {
+                    ModuleBase::WARNING("ESolver_DirectMin_LCAO", "Unable to seed DirectMin variable from KS solver; using random initialization");
+                }
+            }
         }
+        else
+        {
+            ModuleBase::WARNING_QUIT("ESolver_DirectMin_LCAO", "Only 'ks' is supported for directmin_init_method currently");
+        }
+
+        X = std::move(initial_variable);
     }
 
     // next set up the solver based on directmin_solver, sd, cg, bfgs, etc.

source/source_esolver/esolver_ks_lcao.h

@@ -2,6 +2,7 @@
 #define ESOLVER_KS_LCAO_H
 
 #include "esolver_ks.h"
+#include <stdexcept>
 
 // for adjacent atoms
 #include "source_lcao/record_adj.h"
@@ -62,6 +63,42 @@ class ESolver_KS_LCAO : public ESolver_KS<TK>
 
     void after_all_runners(UnitCell& ucell) override;
 
+    const psi::Psi<TK>& get_wavefunctions() const
+    {
+        if (this->psi == nullptr)
+        {
+            throw std::runtime_error("ESolver_KS_LCAO wavefunctions are not initialized");
+        }
+        return *(this->psi);
+    }
+
+    psi::Psi<TK>& get_wavefunctions()
+    {
+        if (this->psi == nullptr)
+        {
+            throw std::runtime_error("ESolver_KS_LCAO wavefunctions are not initialized");
+        }
+        return *(this->psi);
+    }
+
+    const elecstate::ElecState& get_elecstate() const
+    {
+        if (this->pelec == nullptr)
+        {
+            throw std::runtime_error("ESolver_KS_LCAO elecstate is not initialized");
+        }
+        return *(this->pelec);
+    }
+
+    elecstate::ElecState& get_elecstate()
+    {
+        if (this->pelec == nullptr)
+        {
+            throw std::runtime_error("ESolver_KS_LCAO elecstate is not initialized");
+        }
+        return *(this->pelec);
+    }
+
   protected:
     virtual void before_scf(UnitCell& ucell, const int istep) override;