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

python/pyabacus/src/ModuleDriver/py_driver.cpp

@@ -451,7 +451,7 @@ CalculationResult PyDriver::run(const std::string& input_dir,
         if (cal == "scf" || cal == "relax" || cal == "cell-relax" || cal == "nscf")
         {
             Relax_Driver rl_driver;
-            rl_driver.relax_driver(impl_->p_esolver_, *impl_->ucell_, PARAM.inp, GlobalV::ofs_running);
+            rl_driver.relax_driver(impl_->p_esolver_, *impl_->ucell_, PARAM, GlobalV::ofs_running);
         }
         else if (cal == "get_s")
         {

source/Makefile.Objects

@@ -271,6 +271,7 @@ OBJS_DEEPKS=lcao_deepks.o\
 
 
 OBJS_ELECSTAT=elecstate.o\
+    electronic_occupations.o\
     init_scf.o\
     param_update.o\
     estate_e_terms.o\
@@ -298,6 +299,7 @@ OBJS_ELECSTAT=elecstate.o\
     efield.o\
     gatefield.o\
     potential_new.o\
+    electrostatic_potential.o\
     potential_types.o\
     pot_surchem.o\
     pot_sep.o\
@@ -627,6 +629,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/esolver.h

@@ -6,6 +6,14 @@
 #include "source_cell/unitcell.h"
 
 struct Input_para;
+class Parameter;
+
+namespace elecstate
+{
+struct ElectronicControl;
+struct PotentialWindow;
+struct ElectronicObservables;
+}
 
 namespace ModuleESolver
 {
@@ -46,6 +54,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 elecstate::ElectronicControl&) { return false; }
+    virtual bool electronic_observables(const UnitCell&, const Parameter&, const elecstate::PotentialWindow&,
+                                       elecstate::ElectronicObservables&) const { return false; }
+
     bool conv_esolver = true; // whether esolver is converged
 
     std::string classname;

source/source_esolver/esolver_fp.cpp

@@ -1,4 +1,10 @@
 #include "esolver_fp.h"
+#include "source_estate/electronic_control.h"
+#include "source_estate/electronic_occupations.h"
+#include "source_estate/elecstate_tools.h"
+#include "source_estate/module_pot/electrostatic_potential.h"
+
+#include <stdexcept>
 
 #include "source_basis/module_pw/pw_basis_big.h" // use PW_Basis_Big
 #include "source_base/tool_quit.h"
@@ -24,6 +30,34 @@
 #include "source_base/module_parallel/para_bridge.h"
 namespace ModuleESolver
 {
+bool ESolver_FP::apply_electronic_control(Parameter& parameters, const elecstate::ElectronicControl& control)
+{
+    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);
+    elecstate::update_electron_count(*pelec, control.nelect, parameters.inp.nspin,
+                                   parameters.inp.nupdown, parameters.globalv.two_fermi);
+    return true;
+}
+
+bool ESolver_FP::electronic_observables(const UnitCell& ucell, const Parameter& parameters,
+                                       const elecstate::PotentialWindow& window,
+                                       elecstate::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 = elecstate::checked_occupation_electron_count(pelec->wg, parameters.inp.nelec,
+                                                               npool, pw_rhod->poolnproc);
+    result.fermi_ry = elecstate::electronic_chemical_potential(pelec->eferm, parameters.globalv.two_fermi);
+    if (window.axis >= 0)
+    {
+        result.reference_ry = elecstate::electrostatic_reference(ucell, *pw_rhod, *pelec->charge,
+                                                                *pelec->pot, solvent, parameters.inp, window);
+    }
+    return true;
+}
+
 
 ESolver_FP::ESolver_FP()
 {

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 elecstate::ElectronicControl& control) override;
+    bool electronic_observables(const UnitCell& ucell, const Parameter& parameters,
+                               const elecstate::PotentialWindow& window, elecstate::ElectronicObservables& result) const override;
+
   protected:
     virtual void before_scf(UnitCell& ucell, const int istep);
 

source/source_esolver/esolver_ks.cpp

@@ -179,12 +179,14 @@ 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");
 
     const int iter_limit = this->scf_iter_limit(istep);
-    const bool converged = this->scf_loop(ucell, istep, iter_limit);
+    const bool converged = this->scf_loop(ucell, istep, iter_limit) && !this->oscillate_esolver;
     this->after_scf(ucell, istep, converged);
     this->conv_esolver = converged;
     ModuleBase::timer::end(this->classname, "runner");

source/source_esolver/esolver_ks_lcao.cpp

@@ -162,6 +162,17 @@ void ESolver_KS_LCAO<TK, TR>::before_all_runners(BaseCell& basecell, const Input
     return;
 }
 
+template <typename TK, typename TR>
+bool ESolver_KS_LCAO<TK, TR>::apply_electronic_control(Parameter& parameters, const elecstate::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

@@ -35,6 +35,8 @@ class ESolver_KS_LCAO : public ESolver_KS
     //! Release the electronic wave functions.
     ~ESolver_KS_LCAO();
 
+    bool apply_electronic_control(Parameter& parameters, const elecstate::ElectronicControl& control) override;
+
     //! Initialize the LCAO basis, electronic state, and potentials before all steps.
     void before_all_runners(BaseCell& basecell, const Input_para& inp) override;
 

source/source_esolver/esolver_ks_pw.cpp

@@ -18,6 +18,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.h"
@@ -159,6 +160,16 @@ 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 elecstate::ElectronicControl& control)
+{
+    const int functional_type = XC_Functional::get_func_type();
+    validate_electronic_step_exx(functional_type, this->general_exx_info_.cal_exx);
+    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 elecstate::ElectronicControl& control) override;
+
   protected:
     virtual void before_scf(UnitCell& ucell, const int istep) override;
 

source/source_estate/CMakeLists.txt

@@ -2,6 +2,7 @@ list(
   APPEND
   objects
   elecstate.cpp
+  electronic_occupations.cpp
   init_scf.cpp
   estate_e_terms.cpp
   elecstate_energy.cpp
@@ -19,6 +20,7 @@ list(
   module_pot/gatefield.cpp
   module_pot/efield.cpp
   module_pot/h_hartree_pw.cpp
+  module_pot/electrostatic_potential.cpp
   module_pot/h_tddft_pw.cpp
   module_pot/pot_xc.cpp
   module_pot/pot_local.cpp

source/source_estate/elecstate.cpp

@@ -16,11 +16,16 @@ const double* ElecState::getRho(int spin) const
 
 void ElecState::init_nelec_spin()
 {
-    this->nelec_spin.resize(PARAM.inp.nspin);
-    if (PARAM.inp.nspin == 2)
+    this->init_nelec_spin(PARAM.inp.nspin, PARAM.inp.nelec, PARAM.inp.nupdown);
+}
+
+void ElecState::init_nelec_spin(int nspin, double nelect, double nupdown)
+{
+    this->nelec_spin.resize(nspin);
+    if (nspin == 2)
     {
-        this->nelec_spin[0] = (PARAM.inp.nelec + PARAM.inp.nupdown) / 2.0;
-        this->nelec_spin[1] = (PARAM.inp.nelec - PARAM.inp.nupdown) / 2.0;
+        this->nelec_spin[0] = (nelect + nupdown) / 2.0;
+        this->nelec_spin[1] = (nelect - nupdown) / 2.0;
     }
 }
 

source/source_estate/elecstate.h

@@ -116,6 +116,7 @@ class ElecState
     // if nupdown is not 0(TWO_EFERMI case),
     // nelec_spin will be fixed and weights will be constrained
     void init_nelec_spin();
+    void init_nelec_spin(int nspin, double nelect, double nupdown);
     // used to record number of electrons per spin index
     // for NSPIN=2, it will record number of spin up and number of spin down
     // for NSPIN=4, it will record total number, magnetization for x, y, z direction

source/source_estate/elecstate_tools.cpp

@@ -1,6 +1,9 @@
 #include "elecstate_tools.h"
 
 #include "occupy.h"
+
+#include <cmath>
+#include <stdexcept>
 #include "source_base/module_parallel/para_band_output.h"
 #include "source_base/module_parallel/para_bridge.h"
 #include "source_base/parallel_reduce.h"
@@ -9,6 +12,25 @@
 
 namespace elecstate
 {
+void update_electron_count(ElecState& state, double nelect, int nspin, double nupdown, bool two_fermi)
+{
+    state.eferm.two_efermi = two_fermi;
+    state.init_nelec_spin(nspin, nelect, nupdown);
+}
+
+double electronic_chemical_potential(const Efermi& fermi, bool two_fermi)
+{
+    double chemical_potential = fermi.ef;
+    if (two_fermi)
+    {
+        // Total-N derivative at fixed N_up-N_down.
+        chemical_potential = 0.5 * (fermi.ef_up + fermi.ef_dw);
+    }
+    if (!std::isfinite(chemical_potential))
+    { throw std::runtime_error("non-finite SCF chemical potential"); }
+    return chemical_potential;
+}
+
 void calEBand(const ModuleBase::matrix& ekb, const ModuleBase::matrix& wg, fenergy& f_en)
 {
     ModuleBase::TITLE("ElecState", "calEBand");

source/source_estate/elecstate_tools.h

@@ -5,6 +5,9 @@
 
 namespace elecstate
 {
+void update_electron_count(ElecState& state, double nelect, int nspin, double nupdown, bool two_fermi);
+double electronic_chemical_potential(const Efermi& fermi, bool two_fermi);
+
 void calEBand(const ModuleBase::matrix& ekb, const ModuleBase::matrix& wg, fenergy& f_en);
 
 void calculate_weights(const ModuleBase::matrix& ekb,

source/source_estate/electronic_control.h

@@ -0,0 +1,27 @@
+#ifndef ABACUS_ELECTRONIC_CONTROL_H
+#define ABACUS_ELECTRONIC_CONTROL_H
+
+namespace elecstate
+{
+// 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_estate/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 elecstate
+{
+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_estate/electronic_occupations.h

@@ -0,0 +1,20 @@
+#ifndef ABACUS_ELECTRONIC_OCCUPATIONS_H
+#define ABACUS_ELECTRONIC_OCCUPATIONS_H
+
+namespace ModuleBase
+{
+class matrix;
+}
+
+namespace elecstate
+{
+// 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