Merge pull request #4492 from ye-luo/restrict-eref-memory

Restrict Eref update history

docs/methods.rst

@@ -1758,6 +1758,41 @@ Batched ``dmc`` driver (experimental)
 - ``spin_mass`` Optional parameter to allow the user to change the rate of spin sampling. If spin sampling is on using ``spinor`` == yes in the electron ParticleSet input,  the spin mass determines the rate
   of spin sampling, resulting in an effective spin timestep :math:`\tau_s = \frac{\tau}{\mu_s}`. The algorithm is described in detail in :cite:`Melton2016-1` and :cite:`Melton2016-2`.
 
+- ``warmupsteps``: These are the steps at the beginning of a DMC run in
+  which the instantaneous population average energy is used to update the trial
+  energy and updates happen at every step. The aim is to rapidly equilibrate the population while avoiding overly large population fluctuations.
+  Unlike VMC, these warmupsteps are included in the requested DMC step count.
+
+.. math::
+
+  E_\text{trial} = E_\text{pop\_avg}+(\ln \texttt{targetwalkers}-\ln N_\text{pop}) / \texttt{timestep}
+
+where :math:`E_\text{pop\_avg}` is the local energy average over the walker population at the current step
+and :math:`N_\text{pop}` is the current walker population size.
+After the warm-up phase, the trial energy is updated as
+
+.. math::
+
+  E_\text{trial} = E_\text{ref}+\texttt{feedback}\cdot(\ln\texttt{targetWalkers}-\ln N_\text{pop})
+
+where :math:`E_\text{ref}` is the :math:`E_\text{pop\_avg}` average over all the post warm-up steps up to the current step. The update frequency is controlled by ``energyUpdateInterval``.
+
+- ``energyUpdateInterval``: Post warm-up, the trial energy is updated every
+  ``energyUpdateInterval`` steps. Default value is 1 (every step).
+
+- ``refEnergy``: The default reference energy is taken from the VMC run
+  that precedes the DMC run. This value is updated to the current mean
+  whenever branching happens.
+
+- ``feedback``: This variable is used to determine how strong to react
+  to population fluctuations when doing population control. Default value is 1. See the
+  equation in ``warmupsteps`` for more details.
+
+- ``refenergy_update_scheme``: choose a scheme for updating :math:`E_\text{ref}` used in calculating :math:`E_\text{trial}`.
+  'unlimited_history' uses unweighted average of :math:`E_\text{pop\_avg}` of all the steps collected post warm-up.
+  'limited_history' uses weighted average of :math:`E_\text{pop\_avg}` of the latest at maximum
+  min(1, int(1.0 / (feedback * tau))) steps collected post warm-up. Default 'unlimited_history'.
+
 .. code-block::
   :caption: The following is an example of a minimal DMC section using the batched ``dmc`` driver
   :name: Listing 48b

src/Estimators/SizeLimitedDataQueue.hpp

@@ -0,0 +1,83 @@
+//////////////////////////////////////////////////////////////////////////////////////
+// This file is distributed under the University of Illinois/NCSA Open Source License.
+// See LICENSE file in top directory for details.
+//
+// Copyright (c) 2023 QMCPACK developers.
+//
+// File developed by: Ye Luo, [email protected], Argonne National Laboratory
+//
+// File created by: Ye Luo, [email protected], Argonne National Laboratory
+//////////////////////////////////////////////////////////////////////////////////////
+// -*- C++ -*-
+#ifndef QMCPLUSPLUS_SIZELIMITEDDATAQUEUE_H
+#define QMCPLUSPLUS_SIZELIMITEDDATAQUEUE_H
+
+#include <deque>
+#include <array>
+#include <cassert>
+
+namespace qmcplusplus
+{
+
+/** collect data with a history limit.
+ * data stored in std::deque<std::array<T, NUM_FIELDS>>
+ */
+template<typename T, size_t NUM_FIELDS>
+class SizeLimitedDataQueue
+{
+public:
+  struct HistoryElement
+  {
+    T weight;
+    std::array<T, NUM_FIELDS> properties;
+  };
+
+  using value_type = HistoryElement;
+
+  SizeLimitedDataQueue(size_t size_limit) : size_limit_(size_limit) {}
+
+  /// add a new record
+  void push(const value_type& val)
+  {
+    if (data.size() == size_limit_)
+      data.pop_front();
+    assert(data.size() < size_limit_);
+    data.push_back(val);
+  }
+
+  /// add a new record
+  void push(value_type&& val)
+  {
+    if (data.size() == size_limit_)
+      data.pop_front();
+    assert(data.size() < size_limit_);
+    data.push_back(val);
+  }
+
+  /// return weighted average
+  auto weighted_avg() const
+  {
+    std::array<T, NUM_FIELDS> avg;
+    std::fill(avg.begin(), avg.end(), T(0));
+    T weight_sum = 0;
+    for (auto& element : data)
+    {
+      weight_sum += element.weight;
+      for (size_t i = 0; i < NUM_FIELDS; i++)
+        avg[i] += element.properties[i] * element.weight;
+    }
+    for (size_t i = 0; i < NUM_FIELDS; i++)
+      avg[i] /= weight_sum;
+    return avg;
+  }
+
+  /// return the number of records
+  auto size() const { return data.size(); }
+
+private:
+  std::deque<value_type> data;
+  const size_t size_limit_;
+};
+
+} // namespace qmcplusplus
+#endif

src/Estimators/tests/CMakeLists.txt

@@ -33,6 +33,7 @@ set(SRCS
     test_InputSection.cpp
     test_EstimatorManagerInput.cpp
     test_ScalarEstimatorInputs.cpp
+    test_SizeLimitedDataQueue.cpp
     test_MomentumDistribution.cpp
     test_OneBodyDensityMatricesInput.cpp
     test_OneBodyDensityMatrices.cpp

src/Estimators/tests/test_SizeLimitedDataQueue.cpp

@@ -0,0 +1,49 @@
+//////////////////////////////////////////////////////////////////////////////////////
+// This file is distributed under the University of Illinois/NCSA Open Source License.
+// See LICENSE file in top directory for details.
+//
+// Copyright (c) 2023 QMCPACK developers.
+//
+// File developed by: Ye Luo, [email protected], Argonne National Laboratory
+//
+// File created by: Ye Luo, [email protected], Argonne National Laboratory
+//////////////////////////////////////////////////////////////////////////////////////
+
+#include "catch.hpp"
+#include "SizeLimitedDataQueue.hpp"
+
+namespace qmcplusplus
+{
+
+TEST_CASE("SizeLimitedDataQueue", "[estimators]")
+{
+  SizeLimitedDataQueue<double, 1> weight_and_energy(3);
+  CHECK(weight_and_energy.size() == 0);
+  {
+    weight_and_energy.push({1.0, {2.0}});
+    CHECK(weight_and_energy.size() == 1);
+    auto avg = weight_and_energy.weighted_avg();
+    CHECK(Approx(avg[0]) == 2.0);
+  }
+  {
+    weight_and_energy.push({3.0, {1.0}});
+    CHECK(weight_and_energy.size() == 2);
+    auto avg = weight_and_energy.weighted_avg();
+    CHECK(Approx(avg[0]) == 1.25);
+  }
+  {
+    SizeLimitedDataQueue<double, 1>::HistoryElement temp{0.5, {3.0}};
+    weight_and_energy.push(std::move(temp));
+    CHECK(weight_and_energy.size() == 3);
+    auto avg = weight_and_energy.weighted_avg();
+    CHECK(Approx(avg[0]) == 1.444444444);
+  }
+  {
+    weight_and_energy.push({0.5, {3.0}});
+    CHECK(weight_and_energy.size() == 3);
+    auto avg = weight_and_energy.weighted_avg();
+    CHECK(Approx(avg[0]) == 1.5);
+  }
+}
+
+} // namespace qmcplusplus

src/QMCDrivers/CMakeLists.txt

@@ -64,6 +64,7 @@ set(QMCDRIVERS
     DMC/DMCFactoryNew.cpp
     DMC/DMCBatched.cpp
     DMC/DMCDriverInput.cpp
+    DMC/DMCRefEnergy.cpp
     DMC/WalkerControlFactory.cpp
     DMC/WalkerReconfiguration.cpp
     DMC/WalkerControl.cpp

src/QMCDrivers/DMC/DMCBatched.cpp

@@ -394,7 +394,12 @@ void DMCBatched::process(xmlNodePtr node)
     // I'd like to do away with this method in DMCBatched.
 
     app_log() << "  Creating the branching engine and walker controler" << std::endl;
-    branch_engine_ = std::make_unique<SFNBranch>(qmcdriver_input_.get_tau(), population_.get_num_global_walkers());
+    const auto refE_update_scheme = dmcdriver_input_.get_refenergy_update_scheme();
+    app_log() << "    Reference energy is updated using the "
+              << (refE_update_scheme == DMCRefEnergyScheme::UNLIMITED_HISTORY ? "unlimited_history" : "limited_history") << " scheme"
+              << std::endl;
+    branch_engine_ =
+        std::make_unique<SFNBranch>(qmcdriver_input_.get_tau(), dmcdriver_input_.get_feedback(), refE_update_scheme);
     branch_engine_->put(node);
 
     walker_controller_ = std::make_unique<WalkerControl>(myComm, Random, dmcdriver_input_.get_reconfiguration());

src/QMCDrivers/DMC/DMCDriverInput.cpp

@@ -19,10 +19,13 @@ void DMCDriverInput::readXML(xmlNodePtr node)
 {
   ParameterSet parameter_set_;
   std::string reconfig_str;
+  std::string refE_update_scheme_str;
   parameter_set_.add(reconfig_str, "reconfiguration", {"no", "yes", "runwhileincorrect"});
   parameter_set_.add(NonLocalMove, "nonlocalmove", {"no", "yes", "v0", "v1", "v3"});
   parameter_set_.add(NonLocalMove, "nonlocalmoves", {"no", "yes", "v0", "v1", "v3"});
   parameter_set_.add(max_age_, "MaxAge");
+  parameter_set_.add(feedback_, "feedback");
+  parameter_set_.add(refE_update_scheme_str, "refenergy_update_scheme", {"unlimited_history", "limited_history"});
 
   // from DMC.cpp put(xmlNodePtr)
   parameter_set_.add(branch_interval_, "branchInterval");
@@ -64,6 +67,11 @@ void DMCDriverInput::readXML(xmlNodePtr node)
 
   if (reserve_ < 1.0)
     throw std::runtime_error("You can only reserve walkers above the target walker count");
+
+  if (refE_update_scheme_str == "unlimited_history")
+    refenergy_update_scheme_ = DMCRefEnergyScheme::UNLIMITED_HISTORY;
+  else
+    refenergy_update_scheme_ = DMCRefEnergyScheme::LIMITED_HISTORY;
 }
 
 std::ostream& operator<<(std::ostream& o_stream, const DMCDriverInput& dmci) { return o_stream; }

src/QMCDrivers/DMC/DMCDriverInput.h

@@ -14,6 +14,7 @@
 
 #include "Configuration.h"
 #include "OhmmsData/ParameterSet.h"
+#include "DMC/DMCRefEnergyScheme.h"
 
 namespace qmcplusplus
 {
@@ -32,6 +33,8 @@ class DMCDriverInput
   bool get_reconfiguration() const { return reconfiguration_; }
   IndexType get_max_age() const { return max_age_; }
   IndexType get_branch_interval() const { return branch_interval_; }
+  double get_feedback() const { return feedback_; }
+  DMCRefEnergyScheme get_refenergy_update_scheme() const { return refenergy_update_scheme_; }
   const std::string& get_non_local_move() const { return NonLocalMove; }
   double get_alpha() const { return alpha_; }
   double get_gamma() const { return gamma_; }
@@ -46,6 +49,10 @@ class DMCDriverInput
    */
   ///Interval between branching
   IndexType branch_interval_ = 1;
+  ///feed back parameter for population control
+  double feedback_ = 1.0;
+  ///input std::string to determine reference energy update scheme
+  DMCRefEnergyScheme refenergy_update_scheme_;
   ///input std::string to determine kill walkers or not
   std::string KillWalker;
   ///input std::string to determine swap walkers among mpi processors

src/QMCDrivers/DMC/DMCRefEnergy.cpp

@@ -0,0 +1,57 @@
+//////////////////////////////////////////////////////////////////////////////////////
+// This file is distributed under the University of Illinois/NCSA Open Source License.
+// See LICENSE file in top directory for details.
+//
+// Copyright (c) 2023 QMCPACK developers.
+//
+// File developed by: Ye Luo, [email protected], Argonne National Laboratory
+//
+// File created by: Ye Luo, [email protected], Argonne National Laboratory
+//////////////////////////////////////////////////////////////////////////////////////
+
+#include <cassert>
+#include "DMCRefEnergy.h"
+
+namespace qmcplusplus
+{
+
+using FullPrecReal = DMCRefEnergy::FullPrecReal;
+
+DMCRefEnergy::DMCRefEnergy(DMCRefEnergyScheme scheme, size_t history_limit)
+    : scheme_(scheme), energy_and_variance_(history_limit)
+{}
+
+std::tuple<FullPrecReal, FullPrecReal> DMCRefEnergy::getEnergyVariance() const
+{
+  if (scheme_ == DMCRefEnergyScheme::LIMITED_HISTORY)
+  {
+    auto avg = energy_and_variance_.weighted_avg();
+    return {avg[ENERGY], avg[VARIANCE]};
+  }
+  else
+    return {energy_hist_.mean(), variance_hist_.mean()};
+}
+
+void DMCRefEnergy::pushWeightEnergyVariance(FullPrecReal weight, FullPrecReal ene, FullPrecReal var)
+{
+  if (scheme_ == DMCRefEnergyScheme::LIMITED_HISTORY)
+    energy_and_variance_.push({weight, {ene, var}});
+  else
+  {
+    energy_hist_(ene);
+    variance_hist_(var);
+  }
+}
+
+size_t DMCRefEnergy::count() const
+{
+  if (scheme_ == DMCRefEnergyScheme::LIMITED_HISTORY)
+    return energy_and_variance_.size();
+  else
+  {
+    assert(energy_hist_.count() == variance_hist_.count());
+    return energy_hist_.count();
+  }
+}
+
+} // namespace qmcplusplus

src/QMCDrivers/DMC/DMCRefEnergy.h

@@ -0,0 +1,64 @@
+//////////////////////////////////////////////////////////////////////////////////////
+// This file is distributed under the University of Illinois/NCSA Open Source License.
+// See LICENSE file in top directory for details.
+//
+// Copyright (c) 2023 QMCPACK developers.
+//
+// File developed by: Ye Luo, [email protected], Argonne National Laboratory
+//
+// File created by: Ye Luo, [email protected], Argonne National Laboratory
+//////////////////////////////////////////////////////////////////////////////////////
+// -*- C++ -*-
+#ifndef QMCPLUSPLUS_DMCREFENERGY_H
+#define QMCPLUSPLUS_DMCREFENERGY_H
+
+#include <tuple>
+#include <Configuration.h>
+#include <Estimators/SizeLimitedDataQueue.hpp>
+#include <Estimators/accumulators.h>
+#include "DMCRefEnergyScheme.h"
+
+namespace qmcplusplus
+{
+/** Handle updating Eref used for calculating the trial energy.
+ */
+class DMCRefEnergy
+{
+public:
+  using FullPrecReal = QMCTraits::FullPrecRealType;
+
+  enum DataLayout
+  {
+    ENERGY = 0,
+    VARIANCE,
+    DATA_SIZE
+  };
+
+private:
+  /// scheme
+  DMCRefEnergyScheme scheme_;
+
+  // legacy scheme data
+  ///a simple accumulator for energy
+  accumulator_set<FullPrecReal> energy_hist_;
+  ///a simple accumulator for variance
+  accumulator_set<FullPrecReal> variance_hist_;
+
+  // limited memory scheme data
+  SizeLimitedDataQueue<FullPrecReal, DataLayout::DATA_SIZE> energy_and_variance_;
+
+public:
+  DMCRefEnergy(DMCRefEnergyScheme scheme, size_t history_limit);
+
+  /// return energy and variance
+  std::tuple<FullPrecReal, FullPrecReal> getEnergyVariance() const;
+
+  /// record weight, energy and variance.
+  void pushWeightEnergyVariance(FullPrecReal weight, FullPrecReal ene, FullPrecReal var);
+
+  /// return record count.
+  size_t count() const;
+};
+
+} // namespace qmcplusplus
+#endif