Enhance RDMFT module with orbital and occupation optimization classes and interface methods for ABACUS data access

Author: Kai Luo

source/source_esolver/esolver_directmin_lcao.cpp

@@ -185,7 +185,7 @@ void ESolver_DirectMin_LCAO::others(UnitCell& ucell, const int istep)
 //template <typename TK, typename TR>
 void ESolver_DirectMin_LCAO::runner(UnitCell& ucell, const int istep)
 {
-    ModuleBase::TITLE("ESolver_DirectMin_LCAO", "runner");
+    ModuleBase::TITLE(this->classname, "runner");
     ModuleBase::timer::tick(this->classname, "runner");
 
     std::cout << "Running DirectMin ESolver with RDMFT" << std::endl;
@@ -399,280 +399,139 @@ void ESolver_DirectMin_LCAO::setup_params(UnitCell& ucell, const Input_para& inp
 
 void ESolver_DirectMin_LCAO::setup_hybrid_optimization(UnitCell& ucell, const Input_para& inp)
 {
-    ModuleBase::TITLE("ESolver_DirectMin_LCAO", "setup_hybrid_optimization");
-    std::cout << "Setting up hybrid two-level optimization" << std::endl;
-
-    // Setup problems for orbital and occupation optimization
-    this->setup_orbital_problem(ucell, inp);
-    this->setup_occupation_problem(ucell, inp);
+    ModuleBase::TITLE(this->classname, "setup_hybrid_optimization");
 
-    // Setup solvers for orbital and occupation optimization
-    this->setup_orbital_solver(ucell, inp);
-    this->setup_occupation_solver(ucell, inp);
-}
-
-void ESolver_DirectMin_LCAO::setup_orbital_problem(UnitCell& ucell, const Input_para& inp)
-{
-    std::cout << "Setting up orbital optimization problem" << std::endl;
+    // Read RDMFT parameters
+    loop_layer_ = GlobalV::PARAM_RDMFT.rdmft_loop_layer;
+    max_iter_occ_ = GlobalV::PARAM_RDMFT.rdmft_max_iter_occ;
+    max_iter_orb_ = GlobalV::PARAM_RDMFT.rdmft_max_iter_orb;
+    occ_opt_method_ = GlobalV::PARAM_RDMFT.rdmft_occ_opt_method;
+    orb_opt_method_ = GlobalV::PARAM_RDMFT.rdmft_orb_opt_method;
 
-    // Create initial occupation vector (from DFT or previous iteration)
-    Vector initial_occupations; // TODO: Initialize from KS calculation
+    // Setup orbital and occupation problems
+    setup_orbital_problem(ucell, inp);
+    setup_occupation_problem(ucell, inp);
 
-    // Create RDMFT_Orbital problem with fixed occupations
-    // this->orbital_prob = new RDMFT_Orbital(ucell, inp, initial_occupations);
+    // Setup orbital and occupation solvers
+    setup_orbital_solver(ucell, inp);
+    setup_occupation_solver(ucell, inp);
 
-    // For now, create a placeholder (will be replaced with actual RDMFT_Orbital)
-    if (inp.directmin_objective == "test") {
-        this->orbital_prob = new TEST(ucell, inp);
-        
-        Manifold* space = this->orbital_prob->GetDomain();
-        if (space != nullptr) {
-            X_orbitals = space->RandominManifold();
-        } else {
-            ModuleBase::WARNING_QUIT("ESolver_DirectMin_LCAO", "Orbital manifold not defined");
-        }
-        X_orbitals.Print("Initial Orbitals");
-    }
+    ModuleBase::GlobalFunc::OUT(GlobalV::ofs_running, "Hybrid optimization setup complete");
 }
 
-void ESolver_DirectMin_LCAO::setup_occupation_problem(UnitCell& ucell, const Input_para& inp)
+void ESolver_DirectMin_LCAO::setup_orbital_problem(UnitCell& ucell, const Input_para& inp)
 {
-    std::cout << "Setting up occupation optimization problem" << std::endl;
+    // Determine dimensions
+    int nbands = GlobalV::NBANDS;
+    int nlocal = GlobalV::NLOCAL;
+    bool gamma_only = GlobalV::GAMMA_ONLY_LOCAL;
 
-    // Create initial orbital matrix (from DFT or previous iteration)
-    Variable initial_orbitals; // TODO: Initialize from KS calculation
+    // Create orbital problem - pass this ESolver as parameter
+    orbital_prob = new ModuleRDMFT::RDMFT_Orbital(nlocal, nbands, gamma_only, static_cast<const void*>(this));
 
-    // Create RDMFT_Occupation problem with fixed orbitals
-    // this->occupation_prob = new RDMFT_Occupation(ucell, inp, initial_orbitals);
+    // Set the domain (Stiefel manifold)
+    ROPTLITE::Manifold* stiefel_manifold = ModuleRDMFT::create_stiefel_manifold(nlocal, nbands, gamma_only);
+    orbital_prob->SetDomain(stiefel_manifold);
 
-    // For now, create a placeholder (will be replaced with actual RDMFT_Occupation)
-    if (inp.directmin_objective == "test") {
-        this->occupation_prob = new TEST(ucell, inp);
-        
-        Manifold* space = this->occupation_prob->GetDomain();
-        if (space != nullptr) {
-            X_occupations = space->RandominManifold();
-        } else {
-            ModuleBase::WARNING_QUIT("ESolver_DirectMin_LCAO", "Occupation manifold not defined");
-        }
-        X_occupations.Print("Initial Occupations");
-    }
-}
-
-void ESolver_DirectMin_LCAO::setup_orbital_solver(UnitCell& ucell, const Input_para& inp)
-{
-    std::cout << "Setting up orbital solver" << std::endl;
+    // Initialize orbital variables
+    X_orbitals = stiefel_manifold->RandominManifold();
 
-    if (this->orbital_solver != nullptr) {
-        delete this->orbital_solver;
-        this->orbital_solver = nullptr;
-    }
-
-    if (this->orb_opt_method_ == "cg") {
-        this->orbital_solver = new RCG(this->orbital_prob, &X_orbitals);
-    } else if (this->orb_opt_method_ == "sd") {
-        this->orbital_solver = new RSD(this->orbital_prob, &X_orbitals);
-    } else if (this->orb_opt_method_ == "bfgs") {
-        this->orbital_solver = new RBFGS(this->orbital_prob, &X_orbitals);
-    } else {
-        ModuleBase::WARNING_QUIT("ESolver_DirectMin_LCAO", "Unknown orbital optimization method: " + this->orb_opt_method_);
-    }
-    
-    // this->orbital_solver->Verbose = FINALRESULT;
-    std::cout << "Orbital solver: " << this->orb_opt_method_ << std::endl;
+    ModuleBase::GlobalFunc::OUT(GlobalV::ofs_running, "Orbital problem setup complete");
+    std::cout << "Orbital manifold: " << (gamma_only ? "Stiefel (real)" : "CStiefel (complex)") << std::endl;
+    std::cout << "Orbital dimensions: " << nlocal << " x " << nbands << std::endl;
 }
 
-void ESolver_DirectMin_LCAO::setup_occupation_solver(UnitCell& ucell, const Input_para& inp)
-{
-    std::cout << "Setting up occupation solver" << std::endl;
-    
-    if (this->occupation_solver != nullptr) {
-        delete this->occupation_solver;
-        this->occupation_solver = nullptr;
-    }
-
-    if (this->occ_opt_method_ == "cg") {
-        this->occupation_solver = new RCG(this->occupation_prob, &X_occupations);
-    } else if (this->occ_opt_method_ == "sd") {
-        this->occupation_solver = new RSD(this->occupation_prob, &X_occupations);
-    } else if (this->occ_opt_method_ == "bfgs") {
-        this->occupation_solver = new RBFGS(this->occupation_prob, &X_occupations);
-    } else {
-        ModuleBase::WARNING_QUIT("ESolver_DirectMin_LCAO", "Unknown occupation optimization method: " + this->occ_opt_method_);
-    }
-    
-    // this->occupation_solver->Verbose = FINALRESULT;
-    std::cout << "Occupation solver: " << this->occ_opt_method_ << std::endl;
-}
-
-void ESolver_DirectMin_LCAO::run_hybrid_optimization(UnitCell& ucell)
+void ESolver_DirectMin_LCAO::setup_occupation_problem(UnitCell& ucell, const Input_para& inp)
 {
-    std::cout << "Starting hybrid two-level optimization" << std::endl;
+    // Number of occupation numbers to optimize
+    int n_occ = GlobalV::NBANDS;
 
-    double energy_old = 1e10;
-    double energy_new = 0.0;
-    int iter = 0;
-    
-    while (iter < this->max_iter && !this->check_convergence()) {
-        iter++;
-        std::cout << "\n=== Hybrid Iteration " << iter << " ===" << std::endl;
-        
-        // Step 1: Optimize orbitals with fixed occupations
-        std::cout << "Optimizing orbitals..." << std::endl;
-        this->optimize_orbitals_fixed_occupations();
-        
-        // Step 2: Optimize occupations with fixed orbitals  
-        std::cout << "Optimizing occupations..." << std::endl;
-        this->optimize_occupations_fixed_orbitals();
-        
-        // Check convergence
-        energy_new = this->cal_energy();
-        this->print_iteration_info(iter, energy_new, 0.0); // TODO: Add gradient norm
-        
-        if (std::abs(energy_new - energy_old) < this->tol_f) {
-            std::cout << "Converged in energy!" << std::endl;
-            break;
-        }
-        energy_old = energy_new;
-    }
-    
-    std::cout << "Hybrid optimization completed after " << iter << " iterations" << std::endl;
-}
-
-void ESolver_DirectMin_LCAO::optimize_orbitals_fixed_occupations()
-{
-    if (this->orbital_solver == nullptr) {
-        ModuleBase::WARNING_QUIT("ESolver_DirectMin_LCAO", "Orbital solver not initialized");
-    }
+    // Create occupation problem - pass this ESolver as parameter
+    occupation_prob = new ModuleRDMFT::RDMFT_Occupation(n_occ, static_cast<const void*>(this));
 
-    // Run orbital optimization for specified number of iterations
-    for (int i = 0; i < this->max_iter_orb_; ++i) {
-        // TODO: Implement single step of orbital optimization
-        // For now, just run the solver
-        break; // Placeholder
-    }
-}
-
-void ESolver_DirectMin_LCAO::optimize_occupations_fixed_orbitals()
-{
-    if (this->occupation_solver == nullptr) {
-        ModuleBase::WARNING_QUIT("ESolver_DirectMin_LCAO", "Occupation solver not initialized");
-    }
+    // Set the domain (Euclidean space)
+    ROPTLITE::Manifold* euclidean_manifold = new ROPTLITE::Euclidean(n_occ);
+    occupation_prob->SetDomain(euclidean_manifold);
 
-    // Run occupation optimization for specified number of iterations
-    for (int i = 0; i < this->max_iter_occ_; ++i) {
-        // TODO: Implement single step of occupation optimization
-        // For now, just run the solver
-        break; // Placeholder
-    }
-}
-
-// ===============================================
-// Route 2: Joint Optimization Functions
-// ===============================================
-
-void ESolver_DirectMin_LCAO::setup_joint_optimization(UnitCell& ucell, const Input_para& inp)
-{
-    ModuleBase::TITLE("ESolver_DirectMin_LCAO", "setup_joint_optimization");
-    std::cout << "Setting up joint optimization on product manifold" << std::endl;
+    // Initialize occupation variables
+    X_occupations = euclidean_manifold->RandominManifold();
 
-    this->setup_multimanifold_problem(ucell, inp);
+    ModuleBase::GlobalFunc::OUT(GlobalV::ofs_running, "Occupation problem setup complete");
+    std::cout << "Occupation manifold: Euclidean" << std::endl;
+    std::cout << "Number of occupation numbers: " << n_occ << std::endl;
 }
 
-void ESolver_DirectMin_LCAO::setup_multimanifold_problem(UnitCell& ucell, const Input_para& inp)
+// Add missing utility methods
+void ESolver_DirectMin_LCAO::print_iteration_info(int iter, double energy, double grad_norm) const
 {
-    std::cout << "Setting up multimanifold problem" << std::endl;
-    
-    if (inp.directmin_objective == "rdmft") {
-        // Create full RDMFT problem on product manifold
-        // this->prob = new RDMFT(ucell, inp);
-        
-        // TODO: Get initial point from RDMFT problem
-        // X = this->prob->GetInitialPoint();
-        
-        std::cout << "RDMFT problem on product manifold created" << std::endl;
-    } else {
-        // Fallback to test problem for development
-        ModuleBase::WARNING_QUIT("ESolver_DirectMin_LCAO", "RDMFT not implemented yet, use test objective");
+    std::cout << "RDMFT Iteration " << iter << ": Energy = " << std::scientific << std::setprecision(8) << energy;
+    if (grad_norm > 0.0) {
+        std::cout << ", |grad| = " << grad_norm;
     }
-}
-
-void ESolver_DirectMin_LCAO::run_joint_optimization(UnitCell& ucell)
-{
-    std::cout << "Starting joint optimization" << std::endl;
+    std::cout << std::endl;
 
-    if (this->solver == nullptr) {
-        ModuleBase::WARNING_QUIT("ESolver_DirectMin_LCAO", "Joint solver not initialized");
+    GlobalV::ofs_running << "RDMFT Iteration " << iter << ": Energy = " << energy;
+    if (grad_norm > 0.0) {
+        GlobalV::ofs_running << ", |grad| = " << grad_norm;
     }
-    
-    // Run the optimization
-    this->solver->Run();
-    
-    // Print final results
-    this->solver->GetXopt().Print("Final Joint Optimization Result");
-    std::cout << "Joint optimization completed" << std::endl;
+    GlobalV::ofs_running << std::endl;
 }
 
-// ===============================================
-// Helper Functions
-// ===============================================
-
 bool ESolver_DirectMin_LCAO::check_convergence() const
 {
-    // TODO: Implement proper convergence checking
-    // For now, return false to let iteration counter control convergence
+    // Simple convergence check - can be improved
+    // For now, rely on the main optimization loop convergence
     return false;
 }
 
-void ESolver_DirectMin_LCAO::print_iteration_info(int iter, double energy, double grad_norm) const
-{
-    std::cout << "Iteration " << iter << ": Energy = " << std::scientific << energy 
-              << ", Grad norm = " << grad_norm << std::endl;
-    GlobalV::ofs_running << "Iteration " << iter << ": Energy = " << std::scientific << energy 
-                         << ", Grad norm = " << grad_norm << std::endl;
-}
-
 void ESolver_DirectMin_LCAO::cleanup_solvers()
 {
-    // Clean up joint optimization solver
-    if (this->solver != nullptr) {
-        delete this->solver;
-        this->solver = nullptr;
+    if (solver != nullptr) {
+        delete solver;
+        solver = nullptr;
     }
-    
-    // Clean up hybrid optimization solvers
-    if (this->orbital_solver != nullptr) {
-        delete this->orbital_solver;
-        this->orbital_solver = nullptr;
+    if (orbital_solver != nullptr) {
+        delete orbital_solver;
+        orbital_solver = nullptr;
     }
-    
-    if (this->occupation_solver != nullptr) {
-        delete this->occupation_solver;
-        this->occupation_solver = nullptr;
+    if (occupation_solver != nullptr) {
+        delete occupation_solver;
+        occupation_solver = nullptr;
     }
 }
 
 void ESolver_DirectMin_LCAO::cleanup_problems()
 {
-    // Clean up joint optimization problem
-    if (this->prob != nullptr) {
-        delete this->prob;
-        this->prob = nullptr;
+    if (prob != nullptr) {
+        delete prob;
+        prob = nullptr;
     }
-    
-    // Clean up hybrid optimization problems
-    if (this->orbital_prob != nullptr) {
-        delete this->orbital_prob;
-        this->orbital_prob = nullptr;
+    if (orbital_prob != nullptr) {
+        delete orbital_prob;
+        orbital_prob = nullptr;
     }
-    
-    if (this->occupation_prob != nullptr) {
-        delete this->occupation_prob;
-        this->occupation_prob = nullptr;
+    if (occupation_prob != nullptr) {
+        delete occupation_prob;
+        occupation_prob = nullptr;
+    }
+    if (manifold != nullptr) {
+        delete manifold;
+        manifold = nullptr;
     }
 }
 
-// template class ESolver_DirectMin_LCAO<double, double>;
-// template class ESolver_DirectMin_LCAO<std::complex<double>, double>;
-// template class ESolver_DirectMin_LCAO<std::complex<double>, std::complex<double>>;
+// Add interface methods to access ABACUS data structures
+const psi::Psi<double>* ESolver_DirectMin_LCAO::get_psi() const
+{
+    return this->psi;
+}
+
+const elecstate::ElecState* ESolver_DirectMin_LCAO::get_pelec() const
+{
+    return this->pelec;
+}
+
+const hamilt::Hamilt<double>* ESolver_DirectMin_LCAO::get_hamilt() const
+{
+    return this->p_hamilt;
 }