Feature: periodic self-consistent continuum solvation (SCCS) last step: solvent-aware cavity

source/Makefile.Objects

@@ -677,7 +677,8 @@ OBJS_SCCS=sccs_cavity.o\
     sccs_coul_2d.o\
     sccs_moments.o\
     sccs_pw_coulomb.o\
-    sccs_response.o
+    sccs_response.o\
+    sccs_sa.o
 
 OBJS_SYMMETRY=symm_other.o\
     symm_rot_spin.o\

source/source_estate/module_pot/pot_sccs.cpp

@@ -11,6 +11,7 @@
 #include "source_hamilt/module_sccs/sccs_ion_charge.h"
 #include "source_hamilt/module_sccs/sccs_ion_force.h"
 #include "source_hamilt/module_sccs/sccs_response.h"
+#include "source_hamilt/module_sccs/sccs_thread_sum.h"
 #include "source_hamilt/module_sccs/sccs_pw_coulomb.h"
 #include "source_io/module_parameter/input_parameter.h"
 
@@ -53,6 +54,11 @@ void sccs_cfg_input(const Input_para& input,
     config.surface_regularization = input.sccs_surface_eta;
     config.cavity.lowpass_p1 = input.sccs_lowpass_p1;
     config.cavity.lowpass_p2 = input.sccs_lowpass_p2;
+    config.cavity.solvent_aware.solvent_radius = input.sccs_solvent_radius;
+    config.cavity.solvent_aware.radial_scale = input.sccs_radial_scale;
+    config.cavity.solvent_aware.radial_spread = input.sccs_radial_spread;
+    config.cavity.solvent_aware.filling_thr = input.sccs_filling_thr;
+    config.cavity.solvent_aware.filling_spread = input.sccs_filling_spread;
     if (input.assume_isolated == "pcc_0d")
     {
         config.boundary = ModuleSccs::Boundary::Pcc0d;
@@ -132,13 +138,19 @@ PotSccs::PotSccs(const ModulePW::PW_Basis* basis,
     this->dynamic_mode = true;
     sccs_active_ = config_.start_drho <= 0.0 || resume.active;
     restart_potential_ = resume.restart_potential;
+    if (resume.probe_ready)
+    {
+        fixed_sources_.reuse_probe(resume.probe_kernel);
+    }
 }
 
 SccsResume PotSccs::resume_state() const
 {
     SccsResume resume;
     resume.active = sccs_active_;
     resume.restart_potential = restart_potential_;
+    resume.probe_ready = fixed_sources_.probe_ready();
+    resume.probe_kernel = fixed_sources_.probe_kernel();
     return resume;
 }
 
@@ -178,14 +190,14 @@ void PotSccs::cal_v_eff(const Charge* charge, const UnitCell* cell, ModuleBase::
     const std::vector<ModuleBase::Vector3<double>>& grid_positions = fixed_sources_.positions();
     std::vector<double> density(basis.nrxx, 0.0);
     std::vector<double> solute_charge(basis.nrxx);
-    double grid_charges[2] = {0.0, 0.0}; // electrons, ions
-    for (int ir = 0; ir < basis.nrxx; ++ir)
-    {
+    const auto add_point = [&](std::size_t ir, std::array<double, 2>& sums) {
         for (int spin = 0; spin < charge->nspin; ++spin) { density[ir] += charge->rho[spin][ir]; }
         solute_charge[ir] = ions[ir] - density[ir];
-        grid_charges[0] += density[ir];
-        grid_charges[1] += ions[ir];
-    }
+        sums[0] += density[ir];
+        sums[1] += ions[ir];
+    };
+    const std::array<double, 2> sums = ModuleSccs::thread_sums<2>(basis.nrxx, add_point);
+    double grid_charges[2] = {sums[0], sums[1]}; // electrons, ions
     // Pool-reduced, so every rank takes the same decision.
     Parallel_Reduce::reduce_pool(grid_charges, 2);
     const double dv = basis.omega / basis.nxyz;
@@ -221,19 +233,24 @@ void PotSccs::cal_v_eff(const Charge* charge, const UnitCell* cell, ModuleBase::
     if (config_.core_electrons)
     {
         const std::vector<double>& core_density = fixed_sources_.core_density();
+#ifdef _OPENMP
+#pragma omp parallel for schedule(static)
+#endif
         for (int ir = 0; ir < basis.nrxx; ++ir)
         {
             density[ir] += core_density[ir];
         }
     }
+    const std::vector<double>& probe_kernel = fixed_sources_.probe_kernel();
     ModuleSccs::SccsResponse response;
-    ModuleSccs::solve_response(density, solute_charge, config_.cavity, solver_, restart_potential_, basis,
-                               cell->tpiba, *coulomb, response);
+    ModuleSccs::solve_response(density, solute_charge, config_.cavity, probe_kernel, solver_,
+                               restart_potential_, basis, cell->tpiba, *coulomb, response);
     // Only the Coulomb-operator FFTs of the sqrt-CG solve, its main cost; the
     // functional below and the cavity derivatives transform outside this count.
     output_.transforms = coulomb->fft_counts();
     ModuleSccs::FunctionalResult functional;
-    ModuleSccs::eval_functional(solute_charge, response, config_, basis, cell->tpiba, *coulomb, functional);
+    ModuleSccs::eval_functional(solute_charge, response, config_, probe_kernel, basis, cell->tpiba, *coulomb,
+                                functional);
     output_.valid = true;
     output_.iterations = response.polarization.iterations;
     output_.warm_started = response.polarization.warm_started;
@@ -242,6 +259,8 @@ void PotSccs::cal_v_eff(const Charge* charge, const UnitCell* cell, ModuleBase::
     output_.reaction_energy = functional.reaction_energy;
     output_.volume = functional.volume;
     output_.surface = functional.surface;
+    output_.filled = response.solvent_aware.enabled();
+    output_.filled_volume = response.solvent_aware.volume;
     output_.far_field_charge = response.far_field_polarization_charge;
     if (config_.collect_report && config_.boundary != ModuleSccs::Boundary::Periodic)
     {
@@ -250,6 +269,9 @@ void PotSccs::cal_v_eff(const Charge* charge, const UnitCell* cell, ModuleBase::
     electrostatic_rydberg_ = 2.0 * functional.reaction_energy;
     non_electrostatic_rydberg_ = 2.0 * (functional.surface_energy + functional.volume_energy);
     electrostatic_potential_.resize(basis.nrxx);
+#ifdef _OPENMP
+#pragma omp parallel for schedule(static)
+#endif
     for (int ir = 0; ir < basis.nrxx; ++ir)
     {
         electrostatic_potential_[ir] = -2.0 * functional.reaction_potential[ir];
@@ -259,6 +281,8 @@ void PotSccs::cal_v_eff(const Charge* charge, const UnitCell* cell, ModuleBase::
     cavity_potential_ = std::move(functional.cavity_potential);
     epsilon_ = std::move(response.epsilon);
     solute_ = std::move(response.solute);
+    local_solute_ = std::move(response.solvent_aware.local);
+    filled_fraction_ = std::move(response.solvent_aware.filling.fraction);
     restart_potential_ = std::move(response.restart_potential);
     ModuleBase::timer::end("PotSccs", "cal_v_eff");
 }
@@ -274,6 +298,9 @@ void PotSccs::add_sol_force(const UnitCell& cell, ModuleBase::matrix& force) con
     const ModulePW::PW_Basis& basis = *this->rho_basis_;
     const std::vector<unitcell::AtomData> atoms = unitcell::get_atom_data(cell.atoms, cell.ntype, cell.lat0);
     std::vector<double> reaction(basis.nrxx);
+#ifdef _OPENMP
+#pragma omp parallel for schedule(static)
+#endif
     for (int ir = 0; ir < basis.nrxx; ++ir)
     {
         reaction[ir] = -0.5 * electrostatic_potential_[ir];
@@ -349,5 +376,18 @@ void PotSccs::add_sol_fields(std::vector<SolventGridField>& fields) const
     cavity.name = "cavity";
     cavity.values = solute_;
     fields.push_back(cavity);
+    // output_.filled comes from FilledCavity::enabled(), set on every rank.
+    if (!output_.filled)
+    {
+        return;
+    }
+    SolventGridField local;
+    local.name = "cavity_local";
+    local.values = local_solute_;
+    fields.push_back(local);
+    SolventGridField fraction;
+    fraction.name = "filled_fraction";
+    fraction.values = filled_fraction_;
+    fields.push_back(fraction);
 }
 } // namespace elecstate

source/source_estate/module_pot/pot_sccs.h

@@ -30,13 +30,16 @@ std::string check_sccs_cell(const ModuleSccs::SccsConfig& config, const UnitCell
 std::string check_charge(const ModuleSccs::SccsConfig& config, const UnitCell& cell, double electron_count);
 
 // State carried from the component of the previous ionic step, which is
-// rebuilt every ionic step: an activation already reached, and the last
-// unshifted sqrt-CG solution as a warm start. The solver keeps the warm start
-// only when it lowers the initial charge residual, so moved atoms are safe.
+// rebuilt every ionic step: an activation already reached, the last
+// unshifted sqrt-CG solution as a warm start, and the solvent-aware probe
+// kernel of the fixed cell and grid. The solver keeps the warm start only when
+// it lowers the initial charge residual, so moved atoms are safe.
 struct SccsResume
 {
     bool active = false;
     std::vector<double> restart_potential;
+    bool probe_ready = false; // same on every pool rank
+    std::vector<double> probe_kernel; // local G vectors
 };
 
 class PotSccs : public PotBase
@@ -56,7 +59,9 @@ class PotSccs : public PotBase
     void add_sol_force(const UnitCell& cell, ModuleBase::matrix& force) const override;
     void sol_energy(double& electrostatic, double& non_electrostatic) const override;
     const std::vector<double>* sol_elec_pot() const override;
-    // eps and the cavity s of the last active evaluation; none before the start.
+    // eps and the cavity s of the last active evaluation, plus the local s
+    // (cavity_local) and the probe solute fraction c (filled_fraction) of a
+    // solvent-aware cavity; none before the start.
     void add_sol_fields(std::vector<SolventGridField>& fields) const override;
 
     // Delayed start: called once per electronic iteration while SCCS waits,
@@ -82,6 +87,8 @@ class PotSccs : public PotBase
     std::vector<double> cavity_potential_;
     std::vector<double> epsilon_;
     std::vector<double> solute_;
+    std::vector<double> local_solute_;
+    std::vector<double> filled_fraction_;
 };
 } // namespace elecstate
 #endif

source/source_estate/module_pot/sccs_fixed_src.cpp

@@ -6,6 +6,7 @@
 #include "source_cell/unitcell.h"
 #include "source_hamilt/module_sccs/sccs_ion_charge.h"
 #include "source_hamilt/module_sccs/sccs_parameters.h"
+#include "source_hamilt/module_sccs/sccs_sa.h"
 
 namespace
 {
@@ -36,6 +37,9 @@ struct SccsFixedSources::Data
     std::vector<double> ionic_density;
     std::vector<double> core_density;
     std::vector<ModuleBase::Vector3<double>> positions;
+    std::vector<double> probe_kernel;
+    // Set on every rank; probe_kernel is collective to compute.
+    bool probe_kernel_ready = false;
 };
 SccsFixedSources::SccsFixedSources()
 {
@@ -51,6 +55,14 @@ bool SccsFixedSources::update(const UnitCell& cell,
     ModuleBase::timer::start("SccsFixedSources", "update");
     // Atoms are replicated, so every pool rank takes the same decision.
     const bool reuse = data_->valid && same_atoms(data_->atoms, atoms);
+    // The probe depends only on the cell and grid, which are fixed here.
+    const bool filled = ModuleSccs::uses_sa(config.cavity.solvent_aware);
+    if (filled && !data_->probe_kernel_ready)
+    {
+        data_->probe_kernel = ModuleSccs::sa_probe_kernel(basis, cell.latvec, cell.lat0,
+                                                          config.cavity.solvent_aware);
+        data_->probe_kernel_ready = true;
+    }
     if (!reuse)
     {
         ModulePW::grid_positions(basis, cell.latvec, cell.lat0, data_->positions);
@@ -69,4 +81,11 @@ bool SccsFixedSources::update(const UnitCell& cell,
 const std::vector<double>& SccsFixedSources::ionic_density() const { return data_->ionic_density; }
 const std::vector<double>& SccsFixedSources::core_density() const { return data_->core_density; }
 const std::vector<ModuleBase::Vector3<double>>& SccsFixedSources::positions() const { return data_->positions; }
+const std::vector<double>& SccsFixedSources::probe_kernel() const { return data_->probe_kernel; }
+bool SccsFixedSources::probe_ready() const { return data_->probe_kernel_ready; }
+void SccsFixedSources::reuse_probe(const std::vector<double>& kernel)
+{
+    data_->probe_kernel = kernel;
+    data_->probe_kernel_ready = true;
+}
 } // namespace elecstate

source/source_estate/module_pot/sccs_fixed_src.h

@@ -11,10 +11,12 @@ namespace unitcell { struct AtomData; }
 namespace ModuleSccs { struct SccsConfig; }
 namespace elecstate
 {
-// Owns ion/core Gaussian densities and Cartesian grid coordinates. The cell,
+// Owns ion/core Gaussian densities, Cartesian grid coordinates and the
+// solvent-aware probe kernel (empty without the filling). The cell,
 // grid and SCCS configuration are fixed for one PotSccs, which is rebuilt every
 // ionic step; update rebuilds only for different atoms and returns true when
-// the previous sources were reused.
+// the previous sources were reused. The probe depends only on the fixed cell
+// and grid, so a kernel of the previous ionic step can be handed over.
 class SccsFixedSources
 {
 public:
@@ -28,6 +30,12 @@ class SccsFixedSources
     const std::vector<double>& ionic_density() const;
     const std::vector<double>& core_density() const;
     const std::vector<ModuleBase::Vector3<double>>& positions() const;
+    const std::vector<double>& probe_kernel() const;
+    // Same value on every pool rank; the kernel itself is empty on a rank
+    // without local G vectors.
+    bool probe_ready() const;
+    // Called with the same ready flag on every pool rank, before update.
+    void reuse_probe(const std::vector<double>& kernel);
 private:
     struct Data;
     std::unique_ptr<Data> data_;

source/source_estate/module_pot/sccs_output.cpp

@@ -137,8 +137,14 @@ void write_sccs_out(std::ostream& output, const SccsOutput& result, int level,
                                                 "Q_EXPECTED/e", real_text(result.expected_charge, value_digits)};
         write_row(output, gauss);
     }
-    const std::vector<std::string> cavity = {"SCCS_CAVITY", "VOLUME/Bohr3", real_text(result.volume, value_digits),
-                                             "SURFACE/Bohr2", real_text(result.surface, value_digits)};
+    std::vector<std::string> cavity = {"SCCS_CAVITY", "VOLUME/Bohr3", real_text(result.volume, value_digits),
+                                       "SURFACE/Bohr2", real_text(result.surface, value_digits)};
+    if (result.filled)
+    {
+        const std::string filled_text = real_text(result.filled_volume, value_digits);
+        cavity.push_back("V_FILLED/Bohr3");
+        cavity.push_back(filled_text);
+    }
     write_row(output, cavity);
     const std::vector<std::string> transforms = {"SCCS_FFT",
                                                  "R2G_CALLS", int_text(result.transforms.forward_calls),

source/source_estate/module_pot/sccs_output.h

@@ -39,6 +39,8 @@ struct SccsOutput
     double reaction_energy = 0.0;
     double volume = 0.0;
     double surface = 0.0;
+    bool filled = false; // solvent-aware cavity
+    double filled_volume = 0.0;
     double far_field_charge = 0.0;
     double expected_charge = 0.0;
     ModuleSccs::CoulombTransformCounts transforms;

source/source_estate/unittests/test_pot_sccs.cpp

@@ -214,6 +214,19 @@ TEST_F(PotSccsTest, InputMapsCavityModelKeys)
     EXPECT_EQ(config.core_spreads, input.sccs_corespread);
     // The cavity model keys leave the preset dielectric unchanged.
     EXPECT_DOUBLE_EQ(config.cavity.epsilon_bulk, 78.3);
+    EXPECT_FALSE(ModuleSccs::uses_sa(config.cavity.solvent_aware));
+    input.sccs_solvent_radius = 3.0;
+    input.sccs_radial_scale = 1.5;
+    input.sccs_radial_spread = 0.4;
+    input.sccs_filling_thr = 0.6;
+    input.sccs_filling_spread = 0.05;
+    elecstate::sccs_cfg_input(input, config, solver);
+    const ModuleSccs::SolventAwareParameters& filling = config.cavity.solvent_aware;
+    EXPECT_DOUBLE_EQ(filling.solvent_radius, 3.0);
+    EXPECT_DOUBLE_EQ(filling.radial_scale, 1.5);
+    EXPECT_DOUBLE_EQ(filling.radial_spread, 0.4);
+    EXPECT_DOUBLE_EQ(filling.filling_thr, 0.6);
+    EXPECT_DOUBLE_EQ(filling.filling_spread, 0.05);
     EXPECT_DOUBLE_EQ(config.start_drho, 0.0);
     EXPECT_FALSE(config.collect_report);
     input.sccs_start_drho = 1e-3;
@@ -291,6 +304,56 @@ TEST_F(PotSccsTest, DelayedActivationHasOneIrreversibleTransitionAndResumes)
     EXPECT_TRUE(potential->empty());
 }
 
+// The probe kernel of the fixed cell and grid is handed to the next ionic
+// step instead of being rebuilt: a seeded identity kernel (c = s) survives the
+// evaluation unchanged.
+TEST_F(PotSccsTest, ResumeCarriesTheSolventAwareProbeKernel)
+{
+    UnitCell cell;
+    Atom atom;
+    cell.lat0 = length;
+    cell.tpiba = tpiba;
+    cell.omega = basis.omega;
+    cell.ntype = 1;
+    cell.nat = 1;
+    cell.atoms = &atom;
+    atom.na = 1;
+    atom.ncpp.zv = 1.0;
+    atom.tau = {ModuleBase::Vector3<double>(0.25, 0.25, 0.25)};
+    const double density_value = 1.0 / basis.omega;
+    std::vector<double> density(basis.nrxx, density_value);
+    double* channels[] = {density.data()};
+    Charge charge;
+    charge.nspin = 1;
+    charge.rho = channels;
+    Input_para input;
+    input.sccs_epsilon = 5.0;
+    input.sccs_tol_rms = 1e-13;
+    input.sccs_tol_max = 1e-12;
+    input.sccs_solvent_radius = 0.5;
+    ModuleSccs::SccsConfig config;
+    ModuleSccs::PolarizationSolverParameters solver;
+    elecstate::sccs_cfg_input(input, config, solver);
+    elecstate::PotSccs first(&basis, config, solver, 1.0, no_resume);
+    EXPECT_FALSE(first.resume_state().probe_ready);
+    ModuleBase::matrix potential(1, basis.nrxx);
+    first.cal_v_eff(&charge, &cell, potential);
+    const elecstate::SccsResume resume = first.resume_state();
+    EXPECT_TRUE(resume.probe_ready);
+    EXPECT_EQ(resume.probe_kernel.size(), static_cast<std::size_t>(basis.npw));
+    elecstate::PotSccs next(&basis, config, solver, 1.0, resume);
+    potential.zero_out();
+    next.cal_v_eff(&charge, &cell, potential);
+    EXPECT_NEAR(next.get_energy(), first.get_energy(), 1e-10);
+    elecstate::SccsResume seeded = resume;
+    seeded.probe_kernel.assign(basis.npw, 1.0);
+    elecstate::PotSccs identity(&basis, config, solver, 1.0, seeded);
+    potential.zero_out();
+    identity.cal_v_eff(&charge, &cell, potential);
+    const std::vector<double> kept = identity.resume_state().probe_kernel;
+    EXPECT_EQ(kept, seeded.probe_kernel);
+}
+
 TEST_F(PotSccsTest, WaitingSccsContributesNoPotentialEnergyOrForce)
 {
     UnitCell cell;