Add constant-potential electron control to the i-PI socket driver

source/Makefile.Objects

@@ -282,6 +282,7 @@ OBJS_ELECSTAT=elecstate.o\
     efield.o\
     gatefield.o\
     potential_new.o\
+    electrostatic_potential.o\
     potential_types.o\
     pot_surchem.o\
     pot_sep.o\
@@ -321,6 +322,8 @@ OBJS_ELECSTAT_LCAO=elecstate_lcao.o\
 OBJS_ESOLVER=esolver_factory.o\
     esolver_ks.o\
     esolver_fp.o\
+    esolver_fp_electronic.o\
+    electronic_occupations.o\
     esolver_ks_pw.o\
     esolver_ks_pw_tddft.o\
     esolver_ks_lcaopw.o\
@@ -590,6 +593,7 @@ OBJS_RELAXATION=relax_data.o\
     socket_driver.o\
     socket_utils.o\
     socket_handlers.o\
+    socket_constant_potential.o\
     cg_base.o\
     bfgs_basic.o\
     relax_driver.o\

source/source_esolver/CMakeLists.txt

@@ -2,6 +2,8 @@ list(APPEND objects
     esolver_factory.cpp
     esolver_ks.cpp
     esolver_fp.cpp
+    esolver_fp_electronic.cpp
+    electronic_occupations.cpp
     esolver_ks_pw.cpp
     esolver_ks_pw_tddft.cpp
     esolver_ks_lcaopw.cpp

source/source_esolver/electronic_control.h

@@ -0,0 +1,27 @@
+#ifndef ABACUS_ELECTRONIC_CONTROL_H
+#define ABACUS_ELECTRONIC_CONTROL_H
+
+namespace ModuleESolver
+{
+// Physical requests and results; no socket or serialization dependencies.
+struct ElectronicControl
+{
+    double nelect = 0.0;
+};
+
+struct PotentialWindow
+{
+    int axis = -1;
+    double minimum_A = 0.0;
+    double maximum_A = 0.0;
+};
+
+struct ElectronicObservables
+{
+    double nelect = 0.0;
+    double fermi_ry = 0.0;
+    double reference_ry = 0.0;
+};
+}
+
+#endif

source/source_esolver/electronic_occupations.cpp

@@ -0,0 +1,65 @@
+#include "electronic_occupations.h"
+#include "source_base/matrix.h"
+#include "source_base/parallel_reduce.h"
+
+#include <cmath>
+#include <iomanip>
+#include <sstream>
+#include <stdexcept>
+
+namespace ModuleESolver
+{
+double checked_occupation_electron_count(const ModuleBase::matrix& weights,
+                                         double requested_nelect,
+                                         int npool,
+                                         int ranks_per_pool)
+{
+    int invalid = !std::isfinite(requested_nelect) || requested_nelect <= 0.0
+                  || npool <= 0 || ranks_per_pool <= 0;
+    double occupied = 0.0;
+    double compensation = 0.0;
+    for (int ik = 0; ik < weights.nr; ++ik)
+    {
+        for (int ib = 0; ib < weights.nc; ++ib)
+        {
+            const double weight = weights(ik, ib);
+            if (!std::isfinite(weight))
+            {
+                invalid = 1;
+                continue;
+            }
+            // Compensated sum of the local k/spin and band shard. Generalized
+            // smearing weights may be negative; only their total is constrained.
+            const double corrected = weight - compensation;
+            const double next = occupied + corrected;
+            compensation = (next - occupied) - corrected;
+            occupied = next;
+        }
+    }
+    if (!std::isfinite(occupied)) { invalid = 1; }
+    // Synchronize failures before a rank can enter a different collective path.
+    Parallel_Reduce::reduce_all(invalid);
+    if (invalid)
+    {
+        throw std::runtime_error("constant-potential electron count requires finite occupations "
+                                 "and a valid target/occupation partition");
+    }
+
+    // Same distribution as the native band-energy sum: remove the replicated
+    // copies inside each pool before combining KPAR and BNDPAR contributions.
+    Parallel_Reduce::reduce_double_allpool(npool, ranks_per_pool, occupied);
+    const double tolerance = 1.0e-8;
+    int mismatch = !std::isfinite(occupied) || occupied <= 0.0
+                   || std::abs(occupied - requested_nelect) > tolerance;
+    Parallel_Reduce::reduce_all(mismatch);
+    if (mismatch)
+    {
+        std::ostringstream message;
+        message << std::setprecision(17)
+                << "constant-potential occupation electron count mismatch: requested=" << requested_nelect
+                << ", occupied=" << occupied << ", absolute tolerance=" << tolerance;
+        throw std::runtime_error(message.str());
+    }
+    return occupied;
+}
+}

source/source_esolver/electronic_occupations.h

@@ -0,0 +1,20 @@
+#ifndef ABACUS_ELECTRONIC_OCCUPATIONS_H
+#define ABACUS_ELECTRONIC_OCCUPATIONS_H
+
+namespace ModuleBase
+{
+class matrix;
+}
+
+namespace ModuleESolver
+{
+// weights already include k-point weights and spin degeneracy. Rows/bands are
+// partitioned across npool groups and replicated on ranks_per_pool ranks.
+// All ranks must participate; invalid occupations or a mismatch fail together.
+double checked_occupation_electron_count(const ModuleBase::matrix& weights,
+                                         double requested_nelect,
+                                         int npool,
+                                         int ranks_per_pool);
+}
+
+#endif

source/source_esolver/esolver.h

@@ -6,9 +6,13 @@
 #include "source_cell/unitcell.h"
 
 struct Input_para;
+class Parameter;
 
 namespace ModuleESolver
 {
+struct ElectronicControl;
+struct PotentialWindow;
+struct ElectronicObservables;
 class ESolver
 {
   public:
@@ -46,6 +50,11 @@ class ESolver
     //! calcualte stress of given cell
     virtual void cal_stress(BaseCell& cell, ModuleBase::matrix& stress) = 0;
 
+    // Optional first-principles capability. Unsupported solvers return false.
+    virtual bool apply_electronic_control(Parameter&, const ElectronicControl&) { return false; }
+    virtual bool electronic_observables(const UnitCell&, const Parameter&, const PotentialWindow&,
+                                       ElectronicObservables&) const { return false; }
+
     bool conv_esolver = true; // whether esolver is converged
 
     std::string classname;

source/source_esolver/esolver_fp.h

@@ -41,6 +41,10 @@ class ESolver_FP : public ESolver
 
     virtual void after_all_runners(BaseCell& basecell) override;
 
+    bool apply_electronic_control(Parameter& parameters, const ElectronicControl& control) override;
+    bool electronic_observables(const UnitCell& ucell, const Parameter& parameters,
+                               const PotentialWindow& window, ElectronicObservables& result) const override;
+
   protected:
     virtual void before_scf(UnitCell& ucell, const int istep);
 

source/source_esolver/esolver_fp_electronic.cpp

@@ -0,0 +1,106 @@
+#include "esolver_fp.h"
+#include "electronic_control.h"
+#include "electronic_occupations.h"
+#include "source_estate/param_update.h"
+#include "source_estate/module_pot/electrostatic_potential.h"
+#include "source_io/module_parameter/parameter.h"
+#include "source_hamilt/module_xc/xc_functional.h"
+
+#include <cmath>
+#include <stdexcept>
+
+namespace ModuleESolver
+{
+bool ESolver_FP::apply_electronic_control(Parameter& parameters, const ElectronicControl& control)
+{
+    const auto& inp = parameters.inp;
+    if (inp.assume_isolated == "makov-payne")
+    {
+        throw std::runtime_error(
+            "constant-potential electron control is incompatible with the energy-only Makov-Payne correction");
+    }
+    if (inp.esolver_type != "ksdft" || (inp.basis_type != "pw" && inp.basis_type != "lcao")
+        || inp.device != "cpu" || inp.precision != "double" || parameters.globalv.use_uspp
+        || inp.sc_mag_switch || inp.deepks_scf || inp.rpa || inp.rdmft || inp.dft_plus_dmft)
+    {
+        throw std::runtime_error(
+            "constant-potential control requires CPU double-precision PW/LCAO KSDFT with "
+            "norm-conserving pseudopotentials; specialized solver corrections are not supported");
+    }
+    const int functional_type = XC_Functional::get_func_type();
+    if (functional_type < 1 || functional_type > 5)
+    {
+        throw std::runtime_error(
+            "constant-potential control requires a supported KS exchange-correlation functional");
+    }
+    if (inp.dft_plus_u != 0 && inp.dft_plus_u != 1)
+    {
+        throw std::runtime_error(
+            "constant-potential forces require dft_plus_u=1; "
+            "the legacy dft_plus_u=2 force implementation is unavailable");
+    }
+    if (inp.dft_plus_u && (inp.init_occ_mat < 0 || inp.init_occ_mat > 1
+                          || inp.yukawa_potential || parameters.globalv.uramping > 0.0))
+    {
+        throw std::runtime_error(
+            "constant-potential DFT+U requires fixed Hubbard U and self-consistent occupations "
+            "(init_occ_mat=0/1, yukawa_potential=false, uramping disabled)");
+    }
+    if (inp.init_chg == "dm_no_renormalize")
+    { throw std::runtime_error("electron-number control requires charge-density renormalization"); }
+    if (inp.nspin != 1 && inp.nspin != 2)
+    { throw std::runtime_error("constant-potential control supports only nspin=1 or 2"); }
+    if (inp.ocp || inp.nbands <= 0 || !std::isfinite(control.nelect) || control.nelect <= 0.0)
+    { throw std::runtime_error("invalid total electron number, band capacity, or fixed occupation array"); }
+    const bool two_fermi = parameters.globalv.two_fermi;
+    if (two_fermi && inp.nspin != 2)
+    { throw std::runtime_error("two_fermi electron control requires nspin=2"); }
+    const double up = 0.5 * (control.nelect + inp.nupdown);
+    const double down = 0.5 * (control.nelect - inp.nupdown);
+    if ((two_fermi && (!std::isfinite(inp.nupdown) || up <= 0.0 || down <= 0.0
+                      || up > inp.nbands || down > inp.nbands))
+        || control.nelect > 2.0 * inp.nbands)
+    {
+        throw std::runtime_error(
+            "electron number exceeds band capacity or gives invalid spin populations; increase nbands");
+    }
+    if (inp.smearing_method == "fixed"
+        && (std::abs(control.nelect - std::round(control.nelect)) > 1.0e-8
+            || (two_fermi && (std::abs(up-std::round(up)) > 1.0e-8 || std::abs(down-std::round(down)) > 1.0e-8))))
+    { throw std::runtime_error("fractional electron populations require a smearing occupation method"); }
+    if (!pelec || !pelec->charge || !pelec->pot)
+    { throw std::runtime_error("electronic state is not initialized for electron control"); }
+
+    elecstate::ParamUpdater::update_electron_count(control.nelect, parameters);
+    pelec->eferm.two_efermi = two_fermi;
+    pelec->init_nelec_spin();
+    return true;
+}
+
+bool ESolver_FP::electronic_observables(const UnitCell& ucell, const Parameter& parameters,
+                                       const PotentialWindow& window, ElectronicObservables& result) const
+{
+    if (!this->conv_esolver || !pelec || !pelec->charge || !pelec->pot)
+    { throw std::runtime_error("electronic observables require a converged SCF state"); }
+    const int npool = parameters.inp.kpar * parameters.inp.bndpar;
+    result.nelect = checked_occupation_electron_count(pelec->wg, parameters.inp.nelec,
+                                                      npool, pw_rhod->poolnproc);
+    result.fermi_ry = pelec->eferm.ef;
+    if (parameters.globalv.two_fermi)
+    {
+        // The total-N derivative at fixed N_up-N_down is the mean spin chemical potential.
+        result.fermi_ry = 0.5 * (pelec->eferm.ef_up + pelec->eferm.ef_dw);
+    }
+    if (!std::isfinite(result.fermi_ry)) { throw std::runtime_error("non-finite SCF chemical potential"); }
+    if (window.axis >= 0)
+    {
+        const double* fixed = pw_rhod->nrxx > 0 ? pelec->pot->get_fixed_v() : nullptr;
+        const auto* pcc_potential = pelec->pot->pcc_electron_potential();
+        const auto potential = elecstate::electrostatic_potential(ucell, *pw_rhod, *pelec->charge,
+                                                                 fixed, solvent, parameters.inp, pcc_potential);
+        result.reference_ry = elecstate::planar_potential_average(ucell, *pw_rhod, potential, window.axis,
+                                                                 window.minimum_A, window.maximum_A);
+    }
+    return true;
+}
+}

source/source_esolver/esolver_ks.cpp

@@ -177,6 +177,8 @@ void ESolver_KS::runner(BaseCell& basecell, const int istep)
     ModuleBase::TITLE("ESolver_KS", "runner");
     ModuleBase::timer::start(this->classname, "runner");
 
+    this->conv_esolver = false;
+
     // 1) before_scf (electronic iteration loops)
     this->before_scf(ucell, istep);
     ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "INIT SCF");
@@ -243,6 +245,7 @@ void ESolver_KS::runner(BaseCell& basecell, const int istep)
     } // end scf iterations
 
     // 7) after scf
+    conv_esolver = conv_esolver && !this->oscillate_esolver;
     this->after_scf(ucell, istep, conv_esolver);
     this->conv_esolver = conv_esolver;
 

source/source_esolver/esolver_ks_lcao.cpp

@@ -142,6 +142,17 @@ void ESolver_KS_LCAO<TK, TR>::before_all_runners(BaseCell& basecell, const Input
 }
 
 
+template <typename TK, typename TR>
+bool ESolver_KS_LCAO<TK, TR>::apply_electronic_control(Parameter& parameters, const ElectronicControl& control)
+{
+    if (!ESolver_FP::apply_electronic_control(parameters, control)) { return false; }
+#ifdef __EXX
+    if (this->exx_nao.exd) { this->exx_nao.exd->begin_electronic_step(); }
+    if (this->exx_nao.exc) { this->exx_nao.exc->begin_electronic_step(); }
+#endif
+    return true;
+}
+
 template <typename TK, typename TR>
 void ESolver_KS_LCAO<TK, TR>::before_scf(UnitCell& ucell, const int istep)
 {

source/source_esolver/esolver_ks_lcao.h

@@ -32,6 +32,8 @@ class ESolver_KS_LCAO : public ESolver_KS
     ESolver_KS_LCAO();
     ~ESolver_KS_LCAO();
 
+    bool apply_electronic_control(Parameter& parameters, const ElectronicControl& control) override;
+
     void before_all_runners(BaseCell& basecell, const Input_para& inp) override;
 
     double cal_energy() override;

source/source_esolver/esolver_ks_pw.cpp

@@ -17,6 +17,7 @@
 
 #include "source_estate/module_charge/chg_routine.h"  // use charge mixing, mohan add 20251006
 #include "source_estate/setup_estate_pw.h"            // mohan add 20251005
+#include "source_hamilt/module_xc/xc_functional.h"
 #include "source_hamilt/module_xc/general_exx_info.h" // for General_Exx_Info type used via general_exx_info_
 #include "source_io/module_ctrl/ctrl_output_pw.h"     // mohan add 20250927
 #include "source_pw/module_pwdft/dftu_base_io.h"      // append_ion_step_snapshot
@@ -26,6 +27,8 @@
 #include "source_pw/module_pwdft/setup_pot.h"      // mohan add 20250929
 #include "source_pw/module_pwdft/update_cell_pw.h" // mohan add 20250309
 
+#include <stdexcept>
+
 namespace ModuleESolver
 {
 
@@ -135,6 +138,23 @@ void ESolver_KS_PW<T, Device>::before_all_runners(BaseCell& basecell, const Inpu
     this->exx_helper->init(ucell, inp, this->pelec->wg, this->general_exx_info_);
 }
 
+template <typename T, typename Device>
+bool ESolver_KS_PW<T, Device>::apply_electronic_control(Parameter& parameters, const ElectronicControl& control)
+{
+    const int functional_type = XC_Functional::get_func_type();
+    // Separate-loop SCAN0 starts as SCAN, before switching to hybrid meta-GGA.
+    // The upstream PW density path updates tau only for non-hybrid meta-GGA.
+    if (functional_type == 5 || (functional_type == 3 && this->general_exx_info_.cal_exx))
+    {
+        throw std::runtime_error(
+            "constant-potential PW hybrid meta-GGA is unsupported: "
+            "upstream kinetic-energy-density updates do not handle hybrid functionals");
+    }
+    if (!ESolver_FP::apply_electronic_control(parameters, control)) { return false; }
+    this->exx_helper->begin_electronic_step();
+    return true;
+}
+
 template <typename T, typename Device>
 void ESolver_KS_PW<T, Device>::before_scf(UnitCell& ucell, const int istep)
 {

source/source_esolver/esolver_ks_pw.h

@@ -34,6 +34,8 @@ class ESolver_KS_PW : public ESolver_KS
 
     Exx_HelperBase* exx_helper = nullptr;
 
+    bool apply_electronic_control(Parameter& parameters, const ElectronicControl& control) override;
+
   protected:
     virtual void before_scf(UnitCell& ucell, const int istep) override;