Repository: QMCPACK/qmcpack
Pull request #5520: fixes for more complex estimators
Author: login: PDoakORNL
Created: 2025-05-16T17:51:27Z | Merged: 2025-05-17T01:06:06Z

## Proposed changes

`OperatorEstBase::zero` becomes a virtual method, more complex estimators are better off not packing all their data into the share `data_` buffer so overriding to customize zeroing is necessary.

`OperatorEstBase::collect`
no longer has the side effect of zeroing the `OperatorEstBase` objects passed to it.

refactor split of collect `OperatorEstBase::makeBlockAverages` otherwise unchanged.

typo in test program.

## What type(s) of changes does this code introduce?
_Delete the items that do not apply_

- New feature
- Refactoring (no functional changes, no api changes)
- Documentation or build script changes

### Does this introduce a breaking change?

- No

## What systems has this change been tested on?
sdgx2

## Checklist

_Update the following with an [x] where the items apply. If you're unsure about any of them, don't hesitate to ask.  This is
simply a reminder of what we are going to look for before merging your code._

* * [x] I have read the pull request guidance and develop docs
* * [x] This PR is up to date with the current state of 'develop'
* * [x] Code added or changed in the PR has been clang-formatted
* * [x] This PR adds tests to cover any new code, or to catch a bug that is being fixed
* * [x] Documentation has been added (if appropriate)



Diff:
diff --git a/src/Estimators/EstimatorManagerNew.cpp b/src/Estimators/EstimatorManagerNew.cpp
index 59b3355339..bd05173c28 100644
--- a/src/Estimators/EstimatorManagerNew.cpp
+++ b/src/Estimators/EstimatorManagerNew.cpp
@@ -136,7 +136,7 @@ EstimatorManagerNew::~EstimatorManagerNew() = default;
  *
  * \todo this should be in in constructor object shouldn't be reused
  * Warning this is different from some "resets" in the code, it does not clear the object
- * 
+ *
  * The number of estimators and their order can vary from the previous state.
  * reinitialized properties before setting up a new BlockAverage data list.
  *
@@ -286,7 +286,7 @@ void EstimatorManagerNew::collectScalarEstimators(const std::vector<RefVector<Sc
   }
 }
 
-void EstimatorManagerNew::collectOperatorEstimators(const std::vector<RefVector<OperatorEstBase>>& crowd_op_ests)
+void EstimatorManagerNew::collectOperatorEstimators(std::vector<RefVector<OperatorEstBase>>& crowd_op_ests)
 {
   for (int iop = 0; iop < operator_ests_.size(); ++iop)
   {
@@ -294,25 +294,12 @@ void EstimatorManagerNew::collectOperatorEstimators(const std::vector<RefVector<
     for (int icrowd = 0; icrowd < crowd_op_ests.size(); ++icrowd)
       this_op_est_for_all_crowds.emplace_back(crowd_op_ests[icrowd][iop]);
     operator_ests_[iop]->collect(this_op_est_for_all_crowds);
+    operator_ests_[iop]->zero(this_op_est_for_all_crowds);
   }
 }
 
-void EstimatorManagerNew::makeBlockAverages(unsigned long accepts, unsigned long rejects)
+void EstimatorManagerNew::reduceBlockData()
 {
-  // accumulate unsigned long counters over ranks.
-  // Blocks ends are infrequent, two MPI transfers to preserve type
-  // these could be replaced with a singple call MPI_struct_type some packing scheme or even
-  // a pack into and out of an fp type that can be assured to hold the integral type exactly
-  // IMHO they should not be primarily stored in a vector with magic indexes
-  std::vector<unsigned long> accepts_and_rejects(my_comm_->size() * 2, 0);
-  accepts_and_rejects[my_comm_->rank()]                    = accepts;
-  accepts_and_rejects[my_comm_->size() + my_comm_->rank()] = rejects;
-  my_comm_->allreduce(accepts_and_rejects);
-  int64_t total_block_accept =
-      std::accumulate(accepts_and_rejects.begin(), accepts_and_rejects.begin() + my_comm_->size(), int64_t(0));
-  int64_t total_block_reject = std::accumulate(accepts_and_rejects.begin() + my_comm_->size(),
-                                               accepts_and_rejects.begin() + my_comm_->size() * 2, int64_t(0));
-
   //Transfer FullPrecisionRead data
   const size_t n1 = AverageCache.size();
   const size_t n2 = n1 + PropertyCache.size();
@@ -341,6 +328,25 @@ void EstimatorManagerNew::makeBlockAverages(unsigned long accepts, unsigned long
     for (int i = 1; i < PropertyCache.size(); i++)
       PropertyCache[i] *= invTotWgt;
   }
+}
+
+void EstimatorManagerNew::makeBlockAverages(unsigned long accepts, unsigned long rejects)
+{
+  // accumulate unsigned long counters over ranks.
+  // Blocks ends are infrequent, two MPI transfers to preserve type
+  // these could be replaced with a singple call MPI_struct_type some packing scheme or even
+  // a pack into and out of an fp type that can be assured to hold the integral type exactly
+  // IMHO they should not be primarily stored in a vector with magic indexes
+  std::vector<unsigned long> accepts_and_rejects(my_comm_->size() * 2, 0);
+  accepts_and_rejects[my_comm_->rank()]                    = accepts;
+  accepts_and_rejects[my_comm_->size() + my_comm_->rank()] = rejects;
+  my_comm_->allreduce(accepts_and_rejects);
+  int64_t total_block_accept =
+      std::accumulate(accepts_and_rejects.begin(), accepts_and_rejects.begin() + my_comm_->size(), int64_t(0));
+  int64_t total_block_reject = std::accumulate(accepts_and_rejects.begin() + my_comm_->size(),
+                                               accepts_and_rejects.begin() + my_comm_->size() * 2, int64_t(0));
+
+  reduceBlockData();
 
   // now we put the correct accept ratio in
   PropertyCache[acceptRatioInd] = static_cast<FullPrecRealType>(total_block_accept) /
diff --git a/src/Estimators/EstimatorManagerNew.h b/src/Estimators/EstimatorManagerNew.h
index 88af5f7685..108484c412 100644
--- a/src/Estimators/EstimatorManagerNew.h
+++ b/src/Estimators/EstimatorManagerNew.h
@@ -59,7 +59,7 @@ class EstimatorManagerNew
   ///destructor
   ~EstimatorManagerNew();
 
-  /** add a "non" physical operator estimator 
+  /** add a "non" physical operator estimator
    *
    *  this is a dratically reduced version of OperatorBase right now it just supports
    *  what the SpinDensityNew estimator needs
@@ -131,8 +131,9 @@ class EstimatorManagerNew
    *  if the crowd context OperatorEstimator holds a copy of the estimator data structure
    *  or more complex if it just collects for instance a list of writes to locations
    *  in the data structure.
+   *  As a side effect of the method all op_ests passed have there accumulated data zeroed.
    */
-  void collectOperatorEstimators(const std::vector<RefVector<OperatorEstBase>>& op_ests);
+  void collectOperatorEstimators(std::vector<RefVector<OperatorEstBase>>& op_ests);
 
   /** get the average of per-block energy and variance of all the blocks
    * Note: this is not weighted average. It can be the same as weighted average only when block weights are identical.
@@ -184,9 +185,15 @@ class EstimatorManagerNew
   // ///return a pointer to the estimator aname
   // ScalarEstimatorBase* getEstimator(const std::string& a);
 
-  /// collect data and write
+  /** reduce accepts and rejects across all ranks
+   *  call reduceBlockData
+   *  add block averages to the energy and var accumulators
+   */
   void makeBlockAverages(unsigned long accept, unsigned long reject);
 
+  /** Does the mpi reduction over the Property and Average caches */
+  void reduceBlockData();
+
   /// write scalars to scalar.dat and h5
   void writeScalarH5();
 
@@ -255,7 +262,7 @@ class EstimatorManagerNew
   std::vector<ObservableHelper> h5desc;
   /** OperatorEst Observables
    *
-   * since the operator estimators are also a close set at compile time
+   * since the operator estimators are also a closed set at compile time
    * they could be treated just like the inputs.
    * However the idea of a shared interface is much more straight forward for
    * them.
diff --git a/src/Estimators/OperatorEstBase.cpp b/src/Estimators/OperatorEstBase.cpp
index 8d2248541a..cd4b28137d 100644
--- a/src/Estimators/OperatorEstBase.cpp
+++ b/src/Estimators/OperatorEstBase.cpp
@@ -2,7 +2,7 @@
 // This file is distributed under the University of Illinois/NCSA Open Source License.
 // See LICENSE file in top directory for details.
 //
-// Copyright (c) 2022 QMCPACK developers.
+// Copyright (c) 2025 QMCPACK developers.
 //
 // File developed by: Peter Doak, doakpw@ornl.gov, Oak Ridge National Laboratory
 //
@@ -17,6 +17,8 @@
 
 namespace qmcplusplus
 {
+using Real = OperatorEstBase::Real;
+
 OperatorEstBase::OperatorEstBase(DataLocality dl) : data_locality_(dl), walkers_weight_(0) {}
 
 OperatorEstBase::OperatorEstBase(const OperatorEstBase& oth)
@@ -29,7 +31,6 @@ void OperatorEstBase::collect(const RefVector<OperatorEstBase>& type_erased_oper
   {
     std::transform(data_.begin(), data_.end(), crowd_oeb.get_data().begin(), data_.begin(), std::plus<>{});
     walkers_weight_ += crowd_oeb.walkers_weight_;
-    crowd_oeb.zero();
   }
 }
 
@@ -43,9 +44,9 @@ void OperatorEstBase::write(hdf_archive& file)
 {
   if (h5desc_.empty())
     return;
-    // We have to do this to deal with the legacy design that Observables using
-    // collectables in mixed precision were accumulated in float but always written
-    // to hdf5 in double.
+  // We have to do this to deal with the legacy design that Observables using
+  // collectables in mixed precision were accumulated in float but always written
+  // to hdf5 in double.
 #ifdef MIXED_PRECISION
   std::vector<QMCT::FullPrecRealType> expanded_data(data_.size(), 0.0);
   std::copy_n(data_.begin(), data_.size(), expanded_data.begin());
@@ -65,6 +66,12 @@ void OperatorEstBase::packData(PooledData<Real>& buffer) const { buffer.add(data
 
 void OperatorEstBase::unpackData(PooledData<Real>& buffer) { buffer.get(data_.begin(), data_.end()); }
 
+void OperatorEstBase::zero(RefVector<OperatorEstBase>& type_erased_operator_estimators) const
+{
+  for (OperatorEstBase& crowd_oeb : type_erased_operator_estimators)
+    crowd_oeb.zero();
+}
+
 void OperatorEstBase::zero()
 {
   if (data_locality_ == DataLocality::rank || data_locality_ == DataLocality::crowd)
diff --git a/src/Estimators/OperatorEstBase.h b/src/Estimators/OperatorEstBase.h
index 6c47400132..defea89c60 100644
--- a/src/Estimators/OperatorEstBase.h
+++ b/src/Estimators/OperatorEstBase.h
@@ -50,7 +50,7 @@ class OperatorEstBase
   /** Shallow copy constructor!
    *  This alows us to keep the default copy constructors for derived classes which
    *  is quite useful to the spawnCrowdClone design.
-   *  Data is likely to be quite large and since the OperatorEstBase design is that the children 
+   *  Data is likely to be quite large and since the OperatorEstBase design is that the children
    *  reduce to the parent it is infact undesirable for them to copy the data the parent has.
    *  Initialization of Data (i.e. call to resize) if any is the responsibility of the derived class.
    */
@@ -59,7 +59,7 @@ class OperatorEstBase
   virtual ~OperatorEstBase() = default;
 
   /** Accumulate whatever it is you are accumulating with respect to walkers
-   * 
+   *
    *  This method is assumed to be called from the crowd context
    *  It provides parallelism with respect to computational effort of the estimator
    *  without causing a global sync.
@@ -81,7 +81,20 @@ class OperatorEstBase
    *
    *  This is assumed to be called from only from one thread per crowds->rank
    *  reduction. Implied is this is during a global sync or there is a guarantee
-   *  that the crowd operator estimators accumulation data is not being written to.
+   *  that the crowd operator estimators accumulation data is not
+   *  being written to.
+   *
+   *  It is assumed by derived classes and it is necessary to support
+   *  derived type data locality schemes that the OperatorEstBase
+   *  derived types in the refvector match.
+   *
+   *  The input operators are not zeroed after collect is called,
+   *  the owner of the operators must handle the accumulated state explicitly.
+   *
+   *  A side effect is walker_weights_ are collected
+   *  as well, so if this is not called from an override the
+   *  walker_weights_ must be collected there.  If it is called they
+   *  must not be collected there.
    *
    *  There could be concurrent operations inside the scope of the collect call.
    */
@@ -130,14 +143,25 @@ class OperatorEstBase
 
   /** Write to previously registered observable_helper hdf5 wrapper.
    *
-   *  if you haven't registered Operator Estimator 
+   *  if you haven't registered Operator Estimator
    *  this will do nothing.
    */
   virtual void write(hdf_archive& file);
 
+  /** Calls zero on every OperatorEstBase in refvector
+   *
+   *  like collect this is intended to be called with a refvector
+   *  where the OperatorEstBase derived types are all the same.
+   *  Derived types overriding this can assume this.
+   */
+  virtual void zero(RefVector<OperatorEstBase>& oebs) const;
+
   /** zero data appropriately for the DataLocality
+   *
+   *  Derived classes that don't solely rely on data_ for
+   *  their accumulated data must override this function.
    */
-  void zero();
+  virtual void zero();
 
   /** Return the total walker weight for this block
    */
diff --git a/src/Estimators/tests/test_manager_mpi.cpp b/src/Estimators/tests/test_manager_mpi.cpp
index 6be2e514a5..deeb688e01 100644
--- a/src/Estimators/tests/test_manager_mpi.cpp
+++ b/src/Estimators/tests/test_manager_mpi.cpp
@@ -23,7 +23,7 @@ namespace testing
 {
 bool EstimatorManagerNewTest::testMakeBlockAverages()
 {
-  if (em.my_comm_->rank() == 1)
+  if (em.my_comm_->rank() == 0)
   {
     estimators_[1].scalars[0](3.0);
     estimators_[1].scalars[1](3.0);


Reviews:
user: login: ye-luo
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-16T17:59:56Z
user: login: ye-luo
body: LGTM. requested some comments and explanation.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-16T18:03:53Z
user: login: PDoakORNL
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-16T18:27:09Z
user: login: ye-luo
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-16T18:32:09Z
user: login: PDoakORNL
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-16T19:00:36Z
user: login: ye-luo
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-16T19:08:25Z
user: login: ye-luo
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-16T20:05:40Z
user: login: ye-luo
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-16T23:28:14Z

Review comments:
diff_hunk: @@ -187,6 +188,8 @@ class EstimatorManagerNew
   /// collect data and write
   void makeBlockAverages(unsigned long accept, unsigned long reject);
 
+  void reduceBlockData();
path: src/Estimators/EstimatorManagerNew.h
user: login: ye-luo
body: Document this new function?
created_at: 2025-05-16T17:59:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -81,11 +81,15 @@ class OperatorEstBase
    *
    *  This is assumed to be called from only from one thread per crowds->rank
    *  reduction. Implied is this is during a global sync or there is a guarantee
-   *  that the crowd operator estimators accumulation data is not being written to.
+   *  that the crowd operator estimators accumulation data is not
+   *  being written to.
+   *  The input operators are not zeroed after collect is called,
+   *  the owner of the operators must handle the accumulated state explicitly.
    *
    *  There could be concurrent operations inside the scope of the collect call.
    */
   virtual void collect(const RefVector<OperatorEstBase>& oebs);
+  virtual void zero(RefVector<OperatorEstBase>& oebs);
path: src/Estimators/OperatorEstBase.h
user: login: ye-luo
body: Could you move this next to the other `zero()` function and document it?
created_at: 2025-05-16T18:00:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -23,7 +23,7 @@ namespace testing
 {
 bool EstimatorManagerNewTest::testMakeBlockAverages()
 {
-  if (em.my_comm_->rank() == 1)
+  if (em.my_comm_->rank() == 0)
path: src/Estimators/tests/test_manager_mpi.cpp
user: login: ye-luo
body: Could you explain this change? The change doesn't seem to be complex related but more like a logic error?
created_at: 2025-05-16T18:03:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -23,7 +23,7 @@ namespace testing
 {
 bool EstimatorManagerNewTest::testMakeBlockAverages()
 {
-  if (em.my_comm_->rank() == 1)
+  if (em.my_comm_->rank() == 0)
path: src/Estimators/tests/test_manager_mpi.cpp
user: login: PDoakORNL
body: The test should work for any number of mpi ranks but due to this typo it works for any legal number of ranks except 1.
created_at: 2025-05-16T18:27:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -23,7 +23,7 @@ namespace testing
 {
 bool EstimatorManagerNewTest::testMakeBlockAverages()
 {
-  if (em.my_comm_->rank() == 1)
+  if (em.my_comm_->rank() == 0)
path: src/Estimators/tests/test_manager_mpi.cpp
user: login: ye-luo
body: Aha. Got it.
created_at: 2025-05-16T18:32:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -23,7 +23,7 @@ namespace testing
 {
 bool EstimatorManagerNewTest::testMakeBlockAverages()
 {
-  if (em.my_comm_->rank() == 1)
+  if (em.my_comm_->rank() == 0)
path: src/Estimators/tests/test_manager_mpi.cpp
user: login: PDoakORNL
body: Also in this case in the commit comment by complex I meant complicated estimators. 
created_at: 2025-05-16T19:00:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -23,7 +23,7 @@ namespace testing
 {
 bool EstimatorManagerNewTest::testMakeBlockAverages()
 {
-  if (em.my_comm_->rank() == 1)
+  if (em.my_comm_->rank() == 0)
path: src/Estimators/tests/test_manager_mpi.cpp
user: login: ye-luo
body: I see. Maybe adjust title?
created_at: 2025-05-16T19:08:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -139,7 +148,18 @@ class OperatorEstBase
    */
   virtual void write(hdf_archive& file);
 
+  /** Calls zero on every OperatorEstBase in refvector
+   *
+   *  like collect this is intended to be called with a refvector
+   *  where the OperatorEstBase derived types are all the same.
+   *  Derived types overriding this can assume this.
+   */
+  virtual void zero(RefVector<OperatorEstBase>& oebs);
path: src/Estimators/OperatorEstBase.h
user: login: ye-luo
body: Since it is supposed to only modify oebs, should this one be marked const?
created_at: 2025-05-16T20:05:40Z
author_association: CONTRIBUTOR

PR comments:
user: login: PDoakORNL
created_at: 2025-05-16T19:27:54Z
body: added documentation
author_association: CONTRIBUTOR
user: login: ye-luo
created_at: 2025-05-16T23:48:13Z
body: Test this please
author_association: CONTRIBUTOR