Refactor RDMFT_LCAO: Add UnitCell reference and implement VariableToOccNumWfc method for improved wavefunction handling

Author: Kai Luo

source/source_esolver/directmin_problems/rdmft_lcao.cpp

@@ -5,6 +5,8 @@
 #include "source_base/timer.h"
 
 #include <complex>
+#include <algorithm>
+#include <type_traits>
 #include <stdexcept>
 
 #include "source_base/tool_quit.h"
@@ -15,6 +17,21 @@
 
 using namespace ROPTLITE;
 
+namespace
+{
+template <typename ValueType>
+inline ValueType convert_complex_entry(const std::complex<double>& entry, std::true_type)
+{
+    return static_cast<ValueType>(entry.real());
+}
+
+template <typename ValueType>
+inline ValueType convert_complex_entry(const std::complex<double>& entry, std::false_type)
+{
+    return static_cast<ValueType>(entry);
+}
+} // namespace
+
 namespace ModuleESolver
 {
 
@@ -24,6 +41,8 @@ RDMFT_LCAO<TK, TR>::RDMFT_LCAO(UnitCell& ucell, const Input_para& inp)
     // this->classname = "RDMFT_LCAO";
     // this->basisname = "LCAO";
 
+    this->ucell_ref_ = &ucell;
+
 #ifdef __EXX
     // 1. currently this initialization must be put in constructor rather than `before_all_runners()`
     //  because the latter is not reused by ESolver_LCAO_TDDFT,
@@ -138,9 +157,23 @@ RDMFT_LCAO<TK, TR>::~RDMFT_LCAO()
 template <typename TK, typename TR>
 realdp RDMFT_LCAO<TK, TR>::f(const Variable& x) const
 {
+    // convert Variable x to occupation numbers and wavefunctions
+    ModuleBase::matrix occ_num;
+    psi::Psi<TK> wfc;
+
+    VariableToOccNumWfc(x, occ_num, wfc);
+
+    if (this->ucell_ref_ == nullptr)
+    {
+        ModuleBase::WARNING("RDMFT_LCAO", "UnitCell reference is unavailable before calling update_elec");
+        throw std::runtime_error("RDMFT_LCAO::f requires a valid UnitCell reference");
+    }
+
+    this->rdmft_solver.update_elec(*this->ucell_ref_, occ_num, wfc);
+
     this->rdmft_solver.cal_Energy(1); // 1 means not to calculate forces here
 
-    std::cout << "Called f() in RDMFT_LCAO, Etotal = " << this->rdmft_solver.Etotal << std::endl;
+    // std::cout << "Called f() in RDMFT_LCAO, Etotal = " << this->rdmft_solver.Etotal << std::endl;
     return this->rdmft_solver.Etotal;
 }
 
@@ -173,6 +206,105 @@ Vector& RDMFT_LCAO<TK, TR>::EucHessianEta(const Variable& x, const Vector& etax,
     return *result;
 }
 
+template <typename TK, typename TR>
+void RDMFT_LCAO<TK, TR>::VariableToOccNumWfc(const Variable& x,
+                                             ModuleBase::matrix& occ_num,
+                                             psi::Psi<TK>& wfc) const
+{
+    const int required_components = this->nk_total;
+    const int expected_elements = required_components * 2;
+    if (x.Getnumofelements() != expected_elements)
+    {
+        throw std::runtime_error("RDMFT_LCAO::VariableToOccNumWfc received a variable with unexpected component count");
+    }
+
+    if (occ_num.nr != this->nk_total || occ_num.nc != this->nbands_total)
+    {
+        occ_num.create(this->nk_total, this->nbands_total, false);
+    }
+
+    const int local_basis_count = this->pv.get_row_size();
+    const int local_band_count = this->pv.ncol_bands;
+    const int available_band_slots = this->pv.get_col_size();
+    const int band_limit = std::min(local_band_count, available_band_slots);
+
+    if (wfc.get_nk() != this->nk_total || wfc.get_nbands() != local_band_count || wfc.get_nbasis() != local_basis_count)
+    {
+        wfc.resize(this->nk_total, local_band_count, local_basis_count);
+    }
+
+    // Occupation numbers occupy the second half of the product element.
+    const int occ_offset = required_components;
+    for (int ik = 0; ik < this->nk_total; ++ik)
+    {
+        const Element& occ_elem = x.GetElement(occ_offset + ik);
+        if (occ_elem.Getrow() != this->nbands_total || occ_elem.Getcol() != 1)
+        {
+            throw std::runtime_error("RDMFT_LCAO::VariableToOccNumWfc encountered an occupation component with unexpected shape");
+        }
+
+        const double* occ_values = occ_elem.ObtainReadData();
+        for (int ib = 0; ib < this->nbands_total; ++ib)
+        {
+            occ_num(ik, ib) = occ_values[ib];
+        }
+    }
+
+    // Wavefunction blocks sit in the first half of the product element.
+    for (int ik = 0; ik < this->nk_total; ++ik)
+    {
+        const Element& wf_elem = x.GetElement(ik);
+        const int global_rows = wf_elem.Getrow();
+        const int global_cols = wf_elem.Getcol();
+
+        if (global_rows != this->nbasis_total || global_cols != this->nbands_total)
+        {
+            throw std::runtime_error("RDMFT_LCAO::VariableToOccNumWfc encountered a wavefunction block with unexpected shape");
+        }
+
+        const double* raw_data = wf_elem.ObtainReadData();
+        const bool is_complex = wf_elem.Getiscomplex();
+
+        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 >= global_cols)
+            {
+                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 >= global_rows)
+                {
+                    continue;
+                }
+
+                const int linear_index = basis_global + global_rows * band_global;
+
+                TK value;
+                if (is_complex)
+                {
+                    const auto* complex_data = reinterpret_cast<const std::complex<double>*>(raw_data);
+                    const std::complex<double> entry = complex_data[linear_index];
+                    value = convert_complex_entry<TK>(entry, typename std::is_same<TK, double>::type());
+                }
+                else
+                {
+                    value = static_cast<TK>(raw_data[linear_index]);
+                }
+
+                wfc(ir_local) = value;
+            }
+        }
+    }
+}
+
 
 
 

source/source_esolver/directmin_problems/rdmft_lcao.h

@@ -77,13 +77,18 @@ class RDMFT_LCAO : public Problem, public ESolver_KS_LCAO<TK, TR>
     Stiefel* mani_st = nullptr;
     Euclidean* mani_euc = nullptr;
 
-
-
-    rdmft::RDMFT<TK, TR> rdmft_solver;
+    mutable rdmft::RDMFT<TK, TR> rdmft_solver;
+    UnitCell* ucell_ref_ = nullptr;
 
     ProductManifold* mani = nullptr;
     // const TwoCenterBundle* two_center_bundle = nullptr;
 
+    // convert a Variable X to occupation numbers and wavefunctions
+    void VariableToOccNumWfc(const Variable &x, 
+        ModuleBase::matrix &occ_num, 
+        psi::Psi<TK> &wfc) const;
+
+
 };
 
 } // namespace ModuleESolver

source/source_esolver/esolver_directmin_lcao.cpp

@@ -98,22 +98,15 @@ void ESolver_DirectMin_LCAO::before_all_runners(UnitCell& ucell, const Input_par
     ModuleBase::TITLE(this->classname, "before_all_runners");
     ModuleBase::timer::tick(this->classname, "before_all_runners");
 
-    // first setup the problem, check inside the objective type
-    this->setup_problem(ucell, inp);
-
-
-
-
-
 
     // Call parent class method properly
     // this->ESolver_KS_LCAO<double, double>::before_all_runners(ucell, inp);
 
     // Initialize the DirectMin solver based on the optimization approach
-    // this->initialize(ucell, inp);
+    this->initialize(ucell, inp);
 
-    // Setup the optimization problem and solver
-    // this->setup_problem(ucell, inp);
+    // first setup the problem, check inside the objective type
+    this->setup_problem(ucell, inp);
     this->setup_solver(ucell, inp);
     this->setup_params(ucell, inp);