From 45d4d637e4927f8e7eeb2f8fa93871dbfc961e16 Mon Sep 17 00:00:00 2001
From: Jaron Krogel <krogeljt@ornl.gov>
Date: Wed, 15 Sep 2021 11:04:47 -0400
Subject: [PATCH 1/6] add momentum distribution

---
 src/Estimators/CMakeLists.txt           |   4 +-
 src/Estimators/EstimatorManagerNew.cpp  |  31 ++-
 src/Estimators/MomentumDistribution.cpp | 329 ++++++++++++++++++++++++
 src/Estimators/MomentumDistribution.h   | 119 +++++++++
 4 files changed, 480 insertions(+), 3 deletions(-)
 create mode 100644 src/Estimators/MomentumDistribution.cpp
 create mode 100644 src/Estimators/MomentumDistribution.h

diff --git a/src/Estimators/CMakeLists.txt b/src/Estimators/CMakeLists.txt
index 760d597bee..ca025cce76 100644
--- a/src/Estimators/CMakeLists.txt
+++ b/src/Estimators/CMakeLists.txt
@@ -22,7 +22,9 @@ set(QMCEST_SRC
     EstimatorManagerCrowd.cpp
     CollectablesEstimator.cpp
     OperatorEstBase.cpp
-    SpinDensityNew.cpp)
+    SpinDensityNew.cpp
+    MomentumDistribution.cpp
+    )
 
 ####################################
 # create libqmcestimators
diff --git a/src/Estimators/EstimatorManagerNew.cpp b/src/Estimators/EstimatorManagerNew.cpp
index 510df22800..9ba7ea87e3 100644
--- a/src/Estimators/EstimatorManagerNew.cpp
+++ b/src/Estimators/EstimatorManagerNew.cpp
@@ -15,6 +15,7 @@
 
 #include "EstimatorManagerNew.h"
 #include "SpinDensityNew.h"
+#include "MomentumDistribution.h"
 #include "QMCHamiltonians/QMCHamiltonian.h"
 #include "Message/Communicate.h"
 #include "Message/CommOperators.h"
@@ -329,16 +330,19 @@ EstimatorManagerNew::EstimatorType* EstimatorManagerNew::getEstimator(const std:
 
 bool EstimatorManagerNew::put(QMCHamiltonian& H, const ParticleSet& pset, xmlNodePtr cur)
 {
-  std::vector<std::string> extra;
+  std::vector<std::string> extra_types;
+  std::vector<std::string> extra_names;
   cur = cur->children;
   while (cur != NULL)
   {
     std::string cname((const char*)(cur->name));
     if (cname == "estimator")
     {
+      std::string est_type("none");
       std::string est_name(MainEstimatorName);
       std::string use_hdf5("yes");
       OhmmsAttributeSet hAttrib;
+      hAttrib.add(est_type, "type");
       hAttrib.add(est_name, "name");
       hAttrib.add(use_hdf5, "hdf5");
       hAttrib.put(cur);
@@ -378,8 +382,20 @@ bool EstimatorManagerNew::put(QMCHamiltonian& H, const ParticleSet& pset, xmlNod
           operator_ests_.emplace_back(
               std::make_unique<SpinDensityNew>(std::move(spdi), pset.Lattice, pset.mySpecies, dl));
       }
+      else if (est_type == "MomentumDistribution")
+      {
+        MomentumDistributionInput mdi;
+        mdi.readXML(cur);
+        DataLocality dl = DataLocality::crowd;
+        operator_ests_.emplace_back(
+          std::make_unique<MomentumDistribution>(std::move(mdi), 
+            pset.getTotalNum(), pset.getTwist(), pset.Lattice, dl));
+      }
       else
-        extra.push_back(est_name);
+      {
+        extra_types.push_back(est_type);
+        extra_names.push_back(est_name);
+      }
     }
     cur = cur->next;
   }
@@ -395,6 +411,17 @@ bool EstimatorManagerNew::put(QMCHamiltonian& H, const ParticleSet& pset, xmlNod
     app_log() << "  Using CollectablesEstimator for collectables, e.g. sk, gofr, density " << std::endl;
     Collectables = new CollectablesEstimator(H);
   }
+  // Unrecognized estimators are not allowed
+  if (!extra_types.empty())
+  {
+    app_log() << "\nUnrecognized estimators in input:" << std::endl;
+    for (int i=0; i<extra_types.size(); i++)
+    {
+      app_log() << "  type: "<<extra_types[i]<<"     name: "<<extra_names[i]<<std::endl;
+    }
+    app_log() << std::endl;
+    throw UniformCommunicateError("Unrecognized estimators encountered in input.  See log message for more details.");
+  }
   return true;
 }
 
diff --git a/src/Estimators/MomentumDistribution.cpp b/src/Estimators/MomentumDistribution.cpp
new file mode 100644
index 0000000000..0a056efdcf
--- /dev/null
+++ b/src/Estimators/MomentumDistribution.cpp
@@ -0,0 +1,329 @@
+//////////////////////////////////////////////////////////////////////////////////////
+// This file is distributed under the University of Illinois/NCSA Open Source License.
+// See LICENSE file in top directory for details.
+//
+// Copyright (c) 2020 QMCPACK developers.
+//
+// File developed by: Jaron T. Krogel, krogeljt@ornl.gov, Oak Ridge National Laboratory
+//
+// File refactored from: MomentumEstimator.cpp
+//////////////////////////////////////////////////////////////////////////////////////
+
+#include "MomentumDistribution.h"
+#include "CPU/e2iphi.h"
+
+#include <iostream>
+#include <numeric>
+
+
+namespace qmcplusplus
+{
+MomentumDistribution::MomentumDistribution(MomentumDistributionInput&& mdi, size_t np, const PosType& twist_in, const LatticeType& lattice, DataLocality dl)
+    : OperatorEstBase(dl), input_(std::move(mdi)), twist(twist_in), Lattice(lattice)
+{
+  data_locality_ = dl;
+  psi_ratios.resize(np);
+
+  myName = input_.get<std::string>("name");
+  M      = input_.get<int>("samples");
+
+  norm_nofK = 1.0 / RealType(M);
+  
+  //maximum k-value in the k-grid in cartesian coordinates
+  auto kmax = input_.get<RealType>("kmax");
+  //maximum k-values in the k-grid along the reciprocal cell axis
+  auto kmax0 = input_.get<RealType>("kmax0");
+  auto kmax1 = input_.get<RealType>("kmax1");
+  auto kmax2 = input_.get<RealType>("kmax2");
+
+  //dims of a grid for generating k points (obtained below)
+  size_t kgrid = 0;
+  // minimal length as 2 x WS radius.
+  RealType min_Length = Lattice.WignerSeitzRadius_G * 4.0 * M_PI;
+  PosType vec_length;
+  //length of reciprocal lattice vector
+  for (int i = 0; i < OHMMS_DIM; i++)
+    vec_length[i] = 2.0 * M_PI * std::sqrt(dot(Lattice.Gv[i], Lattice.Gv[i]));
+  PosType kmaxs = {kmax0,kmax1,kmax2};
+  RealType sum_kmaxs = kmaxs[0] + kmaxs[1] + kmaxs[2];
+  RealType sphere_kmax;
+  bool sphere      = kmax > 0.0 ? true : false;
+  bool directional = sum_kmaxs > 0.0 ? true : false;
+  if (!sphere && !directional)
+  {
+    // default: kmax = 2 x k_F of polarized non-interacting electron system
+    kmax   = 2.0 * std::pow(6.0 * M_PI * M_PI * np / Lattice.Volume, 1.0 / 3);
+    sphere = true;
+  }
+  sphere_kmax = kmax;
+  for (int i = 0; i < OHMMS_DIM; i++)
+  {
+    if (kmaxs[i] > kmax)
+      kmax = kmaxs[i];
+  }
+  kgrid = int(kmax / min_Length) + 1;
+  RealType kgrid_squared[OHMMS_DIM];
+  for (int i = 0; i < OHMMS_DIM; i++)
+    kgrid_squared[i] = kmaxs[i] * kmaxs[i] / vec_length[i] / vec_length[i];
+  RealType kmax_squared = sphere_kmax * sphere_kmax;
+  std::vector<int> kcount0;
+  std::vector<int> kcount1;
+  std::vector<int> kcount2;
+  kcount0.resize(2 * kgrid + 1, 0);
+  kcount1.resize(2 * kgrid + 1, 0);
+  kcount2.resize(2 * kgrid + 1, 0);
+  for (int i = -kgrid; i < (kgrid + 1); i++)
+  {
+    for (int j = -kgrid; j < (kgrid + 1); j++)
+    {
+      for (int k = -kgrid; k < (kgrid + 1); k++)
+      {
+        PosType ikpt, kpt;
+        ikpt[0] = i - twist[0];
+        ikpt[1] = j - twist[1];
+        ikpt[2] = k - twist[2];
+        //convert to Cartesian: note that 2Pi is multiplied
+        kpt               = Lattice.k_cart(ikpt);
+        bool not_recorded = true;
+        // This collects the k-points within the parallelepiped (if enabled)
+        if (directional && ikpt[0] * ikpt[0] <= kgrid_squared[0] && ikpt[1] * ikpt[1] <= kgrid_squared[1] &&
+            ikpt[2] * ikpt[2] <= kgrid_squared[2])
+        {
+          kPoints.push_back(kpt);
+          kcount0[kgrid + i] = 1;
+          kcount1[kgrid + j] = 1;
+          kcount2[kgrid + k] = 1;
+          not_recorded       = false;
+        }
+        // This collects the k-points within a sphere (if enabled and the k-point has not been recorded yet)
+        if (sphere && not_recorded &&
+            kpt[0] * kpt[0] + kpt[1] * kpt[1] + kpt[2] * kpt[2] <=
+                kmax_squared) //if (std::sqrt(kx*kx+ky*ky+kz*kz)<=sphere_kmax)
+        {
+          kPoints.push_back(kpt);
+        }
+      }
+    }
+  }
+  app_log() <<"\n  MomentumDistribution named "<<myName<<"\n";
+  if (sphere && !directional)
+  {
+    app_log() << "    Using all k-space points with (kx^2+ky^2+kz^2)^0.5 < " << sphere_kmax
+              << " for Momentum Distribution." << std::endl;
+    app_log() << "    Total number of k-points for Momentum Distribution is " << kPoints.size() << std::endl;
+  }
+  else if (directional && !sphere)
+  {
+    int sums[3];
+    sums[0] = 0;
+    sums[1] = 0;
+    sums[2] = 0;
+    for (int i = 0; i < 2 * kgrid + 1; i++)
+    {
+      sums[0] += kcount0[i];
+      sums[1] += kcount1[i];
+      sums[2] += kcount2[i];
+    }
+    app_log() << "    Using all k-space points within cut-offs " << kmax0 << ", " << kmax1 << ", " << kmax2
+              << " for Momentum Distribution." << std::endl;
+    app_log() << "    Total number of k-points for Momentum Distribution: " << kPoints.size() << std::endl;
+    app_log() << "      Number of grid points in kmax0 direction: " << sums[0] << std::endl;
+    app_log() << "      Number of grid points in kmax1 direction: " << sums[1] << std::endl;
+    app_log() << "      Number of grid points in kmax2 direction: " << sums[2] << std::endl;
+  }
+  else
+  {
+    int sums[3];
+    sums[0] = 0;
+    sums[1] = 0;
+    sums[2] = 0;
+    for (int i = 0; i < 2 * kgrid + 1; i++)
+    {
+      sums[0] += kcount0[i];
+      sums[1] += kcount1[i];
+      sums[2] += kcount2[i];
+    }
+    app_log() << "    Using all k-space points with (kx^2+ky^2+kz^2)^0.5 < " << sphere_kmax << ", and" << std::endl;
+    app_log() << "    within the cut-offs " << kmax0 << ", " << kmax1 << ", " << kmax2 << " for Momentum Distribution."
+              << std::endl;
+    app_log() << "    Total number of k-points for Momentum Distribution is " << kPoints.size() << std::endl;
+    app_log() << "    The number of k-points within the cut-off region: " << sums[0] * sums[1] * sums[2] << std::endl;
+    app_log() << "      Number of grid points in kmax0 direction: " << sums[0] << std::endl;
+    app_log() << "      Number of grid points in kmax1 direction: " << sums[1] << std::endl;
+    app_log() << "      Number of grid points in kmax2 direction: " << sums[2] << std::endl;
+  }
+  app_log() << "    Number of samples: " << M << std::endl;
+  app_log() << "    My twist is: " << twist[0] << "  " << twist[1] << "  " << twist[2] << "\n\n";
+
+  // resize arrays
+  nofK.resize(kPoints.size());
+  kdotp.resize(kPoints.size());
+  vPos.resize(M);
+  phases.resize(kPoints.size());
+  phases_vPos.resize(M);
+  for (int im = 0; im < M; im++)
+    phases_vPos[im].resize(kPoints.size());
+  psi_ratios_all.resize(M, psi_ratios.size());
+
+  // allocate data storage
+  size_t data_size  = nofK.size();
+  data_             = createLocalData(data_size, data_locality_);
+
+}
+
+
+MomentumDistribution* MomentumDistribution::clone()
+{
+  auto* md = new MomentumDistribution(*this);
+  if (md->data_locality_ == DataLocality::crowd)
+  {
+    app_log()<<"MD::clone dl crowd\n";
+    size_t data_size = data_->size();
+    md->data_        = createLocalData(data_size, data_locality_);
+  }
+  else if (md->data_locality_ == DataLocality::rank)
+  {
+    app_log()<<"MD::clone dl rank\n";
+    assert(data_locality_ == DataLocality::rank);
+    size_t data_size  = 10; // jtk fix
+    md->data_locality_    = DataLocality::queue;
+    md->data_             = createLocalData(data_size, data_locality_);
+  }
+  else
+    app_log()<<"MD::clone dl other\n";
+  
+  return md;
+}
+
+//MomentumDistribution::MomentumDistribution(const MomentumDistribution& md)
+//    : OperatorEstBase(md),
+//      input_(std::move(md.input_))
+//{
+//  if (data_locality_ == DataLocality::crowd)
+//  {
+//    app_log()<<"MD::cons dl crowd\n";
+//    size_t data_size = md.data_->size();
+//    data_            = createLocalData(data_size, data_locality_);
+//  }
+//  else if (data_locality_ == DataLocality::rank)
+//  {
+//    app_log()<<"MD::cons dl rank\n";
+//    assert(md.data_locality_ == DataLocality::rank);
+//    size_t data_size  = 10; // jtk fix
+//    data_locality_    = DataLocality::queue;
+//    data_             = createLocalData(data_size, data_locality_);
+//  }
+//  else
+//    app_log()<<"MD::cons dl other\n";
+//}
+
+void MomentumDistribution::startBlock(int steps)
+{
+  //if (data_locality_ == DataLocality::rank)
+  //{
+  //  app_log()<<"MD::startBlock dl rank\n";
+  //  size_t data_size  = 10; // jtk fix
+  //  data_->reserve(data_size);
+  //  data_->resize(0);
+  //}
+  //else
+  //  app_log()<<"MD::startBlock dl other\n";
+
+}
+
+/** Gets called every step and writes to thread local data.
+ *
+ */
+void MomentumDistribution::accumulate(const RefVector<MCPWalker>& walkers, const RefVector<ParticleSet>& psets)
+{
+  for (int iw = 0; iw < walkers.size(); ++iw)
+  {
+    MCPWalker& walker = walkers[iw];
+    ParticleSet& pset = psets[iw];
+    RealType weight   = walker.Weight;
+
+    const int np = pset.getTotalNum();
+    const int nk = kPoints.size();
+
+    // compute phase factors
+    for (int s = 0; s < M; ++s)
+    {
+      PosType newpos;
+      // JTK: restore this once access to RNG is provided
+      //for (int i = 0; i < OHMMS_DIM; ++i)
+      //  newpos[i] = myRNG();
+      //make it cartesian
+      vPos[s] = Lattice.toCart(newpos);
+      pset.makeVirtualMoves(vPos[s]);
+      // JTK: restore this with addition of RefVector<TrialWaveFunction>
+      //refPsi.evaluateRatiosAlltoOne(P, psi_ratios);
+      for (int i = 0; i < np; ++i)
+        psi_ratios_all[s][i] = psi_ratios[i];
+      
+      for (int ik = 0; ik < nk; ++ik)
+        kdotp[ik] = -dot(kPoints[ik], vPos[s]);
+      eval_e2iphi(nk, kdotp.data(), phases_vPos[s].data(0), phases_vPos[s].data(1));
+    }
+
+    // update n(k)
+    std::fill_n(nofK.begin(), nk, RealType(0));
+    for (int i = 0; i < np; ++i)
+    {
+      for (int ik = 0; ik < nk; ++ik)
+        kdotp[ik] = dot(kPoints[ik], pset.R[i]);
+      eval_e2iphi(nk, kdotp.data(), phases.data(0), phases.data(1));
+      for (int s = 0; s < M; ++s)
+      {
+        const ComplexType one_ratio(psi_ratios_all[s][i]);
+        const RealType ratio_c                 = one_ratio.real();
+        const RealType ratio_s                 = one_ratio.imag();
+        const RealType* restrict phases_c      = phases.data(0);
+        const RealType* restrict phases_s      = phases.data(1);
+        const RealType* restrict phases_vPos_c = phases_vPos[s].data(0);
+        const RealType* restrict phases_vPos_s = phases_vPos[s].data(1);
+        RealType* restrict nofK_here           = nofK.data();
+#pragma omp simd aligned(nofK_here, phases_c, phases_s, phases_vPos_c, phases_vPos_s : QMC_SIMD_ALIGNMENT)
+        for (int ik = 0; ik < nk; ++ik)
+          nofK_here[ik] += (phases_c[ik] * phases_vPos_c[ik] - phases_s[ik] * phases_vPos_s[ik]) * ratio_c -
+            (phases_s[ik] * phases_vPos_c[ik] + phases_c[ik] * phases_vPos_s[ik]) * ratio_s;
+      }
+    }
+
+    // accumulate data
+    for (int ik = 0; ik < nofK.size(); ++ik)
+      (*data_)[ik] += weight * nofK[ik];
+
+  }
+};
+
+
+void MomentumDistribution::collect(const RefVector<OperatorEstBase>& type_erased_operator_estimators)
+{
+  if (data_locality_ == DataLocality::crowd)
+  {
+    OperatorEstBase::collect(type_erased_operator_estimators);
+  }
+  else
+  {
+    throw std::runtime_error("You cannot call collect on a MomentumDistribution with this DataLocality");
+  }
+}
+
+
+void MomentumDistribution::registerOperatorEstimator(hid_t gid)
+{
+  //descriptor for the data, 1-D data
+  std::vector<int> ng(1);
+  //add nofk
+  ng[0] = nofK.size();
+  h5desc_.emplace_back(std::make_unique<ObservableHelper>("nofk"));
+  auto& h5o = h5desc_.back();
+  //h5o.set_dimensions(ng, myIndex);
+  h5o->set_dimensions(ng, 0); // JTK: doesn't seem right
+  h5o->open(gid);
+  h5o->addProperty(const_cast<std::vector<PosType>&>(kPoints), "kpoints");
+  h5o->addProperty(const_cast<std::vector<int>&>(kWeights), "kweights");    
+}
+
+
+} // namespace qmcplusplus
diff --git a/src/Estimators/MomentumDistribution.h b/src/Estimators/MomentumDistribution.h
new file mode 100644
index 0000000000..5a97cfe1e2
--- /dev/null
+++ b/src/Estimators/MomentumDistribution.h
@@ -0,0 +1,119 @@
+//////////////////////////////////////////////////////////////////////////////////////
+// This file is distributed under the University of Illinois/NCSA Open Source License.
+// See LICENSE file in top directory for details.
+//
+// Copyright (c) 2020 QMCPACK developers.
+//
+// File developed by: Jaron T. Krogel, krogeljt@ornl.gov, Oak Ridge National Laboratory
+//
+// File refactored from: MomentumEstimator.h
+//////////////////////////////////////////////////////////////////////////////////////
+
+#ifndef QMCPLUSPLUS_MOMENTUMDISTRIBUTION_H
+#define QMCPLUSPLUS_MOMENTUMDISTRIBUTION_H
+#include <vector>
+
+#include "Configuration.h"
+#include "OperatorEstBase.h"
+#include "Containers/OhmmsPETE/TinyVector.h"
+
+#include "MomentumDistributionInput.h"
+
+
+namespace qmcplusplus
+{
+/** Class that collects momentum distribution of electrons
+ *  
+ */
+class MomentumDistribution : public OperatorEstBase
+{
+public:
+  using LatticeType = PtclOnLatticeTraits::ParticleLayout_t;
+  using RealType    = QMCTraits::RealType;
+  using ComplexType = QMCTraits::ComplexType;
+  using ValueType   = QMCTraits::ValueType;
+  using PosType     = QMCTraits::PosType;
+
+  //data members set only during construction
+  ///input values
+  MomentumDistributionInput input_;
+  ///number of samples
+  int M;
+  ///reference to the trial wavefunction for ratio evaluations
+  //TrialWaveFunction& refPsi;
+  ///twist angle
+  PosType twist;
+  ///lattice vector
+  LatticeType Lattice;
+  ///normalization factor for n(k)
+  RealType norm_nofK;
+  ///random generator
+  RandomGenerator_t myRNG;
+  ///list of k-points in Cartesian Coordinates
+  std::vector<PosType> kPoints;
+  ///weight of k-points (make use of symmetry)
+  std::vector<int> kWeights;
+
+  /** @ingroup MomentumDistribution mutable data members
+   */
+  ///sample positions
+  std::vector<PosType> vPos;
+  ///wavefunction ratios
+  std::vector<ValueType> psi_ratios;
+  ///wavefunction ratios all samples
+  Matrix<ValueType> psi_ratios_all;
+  ///nofK internal
+  Vector<RealType> kdotp;
+  ///phases
+  VectorSoaContainer<RealType, 2> phases;
+  ///phases of vPos
+  std::vector<VectorSoaContainer<RealType, 2>> phases_vPos;
+  ///nofK
+  aligned_vector<RealType> nofK;
+
+  /** Constructor for MomentumDistributionInput 
+   */
+  MomentumDistribution(MomentumDistributionInput&& mdi, size_t np, const PosType& twist, const LatticeType& lattice, DataLocality dl = DataLocality::crowd);
+
+  //MomentumDistribution(const MomentumDistribution& md);
+
+  /** This allows us to allocate the necessary data for the DataLocality::queue 
+   */
+  void startBlock(int steps) override;
+
+  /** standard interface
+   */
+  MomentumDistribution* clone() override;
+
+  /** accumulate 1 or more walkers of MomentumDistribution samples
+   */
+  void accumulate(const RefVector<MCPWalker>& walkers, const RefVector<ParticleSet>& psets) override;
+
+  /** this allows the EstimatorManagerNew to reduce without needing to know the details
+   *  of MomentumDistribution's data.
+   *
+   *  can use base class default until crowd level MomentumDistribution
+   *  estimators don't have a copy of the density grid.
+   */
+  void collect(const RefVector<OperatorEstBase>& operator_estimators) override;
+
+  /** this allows the EstimatorManagerNew to reduce without needing to know the details
+   *  of MomentumDistribution's data.
+   *
+   *  can use base class default until crowd level MomentumDistribution estimators don't have a copy of the density grid.
+   */
+  //void collect(const OperatorEstBase&  oeb);
+
+  /** this gets us into the hdf5 file
+   *
+   *  Just parroting for now don't fully understand.
+   *, needs to be unraveled and simplified the hdf5 output is another 
+   *  big state big coupling design.
+   */
+  void registerOperatorEstimator(hid_t gid) override;
+
+};
+
+} // namespace qmcplusplus
+
+#endif /* QMCPLUSPLUS_MOMENTUMDISTRIBUTION_H */

From d0354f75acd7f2fb6d0a60432c2cd3e3b70660b2 Mon Sep 17 00:00:00 2001
From: Jaron Krogel <krogeljt@ornl.gov>
Date: Thu, 16 Sep 2021 10:02:29 -0400
Subject: [PATCH 2/6] restore rng/wfn use in MomentumDistribution

---
 src/Estimators/MomentumDistribution.cpp | 18 +++++++++---------
 src/Estimators/MomentumDistribution.h   |  4 +---
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/src/Estimators/MomentumDistribution.cpp b/src/Estimators/MomentumDistribution.cpp
index 0a056efdcf..18ef430b99 100644
--- a/src/Estimators/MomentumDistribution.cpp
+++ b/src/Estimators/MomentumDistribution.cpp
@@ -11,6 +11,7 @@
 
 #include "MomentumDistribution.h"
 #include "CPU/e2iphi.h"
+#include "TrialWaveFunction.h"
 
 #include <iostream>
 #include <numeric>
@@ -234,13 +235,14 @@ void MomentumDistribution::startBlock(int steps)
 /** Gets called every step and writes to thread local data.
  *
  */
-void MomentumDistribution::accumulate(const RefVector<MCPWalker>& walkers, const RefVector<ParticleSet>& psets)
+void MomentumDistribution::accumulate(const RefVector<MCPWalker>& walkers, const RefVector<ParticleSet>& psets, const RefVector<TrialWaveFunction>& wfns, RandomGenerator_t& rng)
 {
   for (int iw = 0; iw < walkers.size(); ++iw)
   {
-    MCPWalker& walker = walkers[iw];
-    ParticleSet& pset = psets[iw];
-    RealType weight   = walker.Weight;
+    MCPWalker& walker      = walkers[iw];
+    ParticleSet& pset      = psets[iw];
+    TrialWaveFunction& psi = wfns[iw];
+    RealType weight        = walker.Weight;
 
     const int np = pset.getTotalNum();
     const int nk = kPoints.size();
@@ -249,14 +251,12 @@ void MomentumDistribution::accumulate(const RefVector<MCPWalker>& walkers, const
     for (int s = 0; s < M; ++s)
     {
       PosType newpos;
-      // JTK: restore this once access to RNG is provided
-      //for (int i = 0; i < OHMMS_DIM; ++i)
-      //  newpos[i] = myRNG();
+      for (int i = 0; i < OHMMS_DIM; ++i)
+        newpos[i] = rng();
       //make it cartesian
       vPos[s] = Lattice.toCart(newpos);
       pset.makeVirtualMoves(vPos[s]);
-      // JTK: restore this with addition of RefVector<TrialWaveFunction>
-      //refPsi.evaluateRatiosAlltoOne(P, psi_ratios);
+      psi.evaluateRatiosAlltoOne(pset, psi_ratios);
       for (int i = 0; i < np; ++i)
         psi_ratios_all[s][i] = psi_ratios[i];
       
diff --git a/src/Estimators/MomentumDistribution.h b/src/Estimators/MomentumDistribution.h
index 5a97cfe1e2..f229dd657b 100644
--- a/src/Estimators/MomentumDistribution.h
+++ b/src/Estimators/MomentumDistribution.h
@@ -39,8 +39,6 @@ class MomentumDistribution : public OperatorEstBase
   MomentumDistributionInput input_;
   ///number of samples
   int M;
-  ///reference to the trial wavefunction for ratio evaluations
-  //TrialWaveFunction& refPsi;
   ///twist angle
   PosType twist;
   ///lattice vector
@@ -87,7 +85,7 @@ class MomentumDistribution : public OperatorEstBase
 
   /** accumulate 1 or more walkers of MomentumDistribution samples
    */
-  void accumulate(const RefVector<MCPWalker>& walkers, const RefVector<ParticleSet>& psets) override;
+  void accumulate(const RefVector<MCPWalker>& walkers, const RefVector<ParticleSet>& psets, const RefVector<TrialWaveFunction>& wfns, RandomGenerator_t& rng) override;
 
   /** this allows the EstimatorManagerNew to reduce without needing to know the details
    *  of MomentumDistribution's data.

From eb3fea93504d45dde190c86b88f82f987c780847 Mon Sep 17 00:00:00 2001
From: Jaron Krogel <krogeljt@ornl.gov>
Date: Thu, 16 Sep 2021 11:01:42 -0400
Subject: [PATCH 3/6] fix incorrect use of unsigned int

---
 src/Estimators/MomentumDistribution.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Estimators/MomentumDistribution.cpp b/src/Estimators/MomentumDistribution.cpp
index 18ef430b99..0cf36e082c 100644
--- a/src/Estimators/MomentumDistribution.cpp
+++ b/src/Estimators/MomentumDistribution.cpp
@@ -38,7 +38,7 @@ MomentumDistribution::MomentumDistribution(MomentumDistributionInput&& mdi, size
   auto kmax2 = input_.get<RealType>("kmax2");
 
   //dims of a grid for generating k points (obtained below)
-  size_t kgrid = 0;
+  int kgrid = 0;
   // minimal length as 2 x WS radius.
   RealType min_Length = Lattice.WignerSeitzRadius_G * 4.0 * M_PI;
   PosType vec_length;

From c41d48cf61e6d279e0328e17b6f4bc84913de3f0 Mon Sep 17 00:00:00 2001
From: Jaron Krogel <krogeljt@ornl.gov>
Date: Tue, 21 Sep 2021 11:46:53 -0400
Subject: [PATCH 4/6] fix n(k) weighting

---
 src/Estimators/MomentumDistribution.cpp | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/Estimators/MomentumDistribution.cpp b/src/Estimators/MomentumDistribution.cpp
index 0cf36e082c..da35a64cb2 100644
--- a/src/Estimators/MomentumDistribution.cpp
+++ b/src/Estimators/MomentumDistribution.cpp
@@ -247,6 +247,10 @@ void MomentumDistribution::accumulate(const RefVector<MCPWalker>& walkers, const
     const int np = pset.getTotalNum();
     const int nk = kPoints.size();
 
+    // accumulate weight 
+    //  (required by all estimators, otherwise inf results)
+    walkers_weight_ += weight;
+
     // compute phase factors
     for (int s = 0; s < M; ++s)
     {
@@ -291,7 +295,7 @@ void MomentumDistribution::accumulate(const RefVector<MCPWalker>& walkers, const
 
     // accumulate data
     for (int ik = 0; ik < nofK.size(); ++ik)
-      (*data_)[ik] += weight * nofK[ik];
+      (*data_)[ik] += weight * nofK[ik] * norm_nofK;
 
   }
 };

From 3a045f4d0c7e7ff0b34698f97b2965bd91eebc4e Mon Sep 17 00:00:00 2001
From: Jaron Krogel <krogeljt@ornl.gov>
Date: Thu, 23 Sep 2021 13:40:08 -0400
Subject: [PATCH 5/6] set members to const

---
 src/Estimators/MomentumDistribution.h | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/src/Estimators/MomentumDistribution.h b/src/Estimators/MomentumDistribution.h
index f229dd657b..33fcab160b 100644
--- a/src/Estimators/MomentumDistribution.h
+++ b/src/Estimators/MomentumDistribution.h
@@ -37,16 +37,14 @@ class MomentumDistribution : public OperatorEstBase
   //data members set only during construction
   ///input values
   MomentumDistributionInput input_;
-  ///number of samples
-  int M;
   ///twist angle
-  PosType twist;
+  const PosType twist;
   ///lattice vector
-  LatticeType Lattice;
+  const LatticeType Lattice;
+  ///number of samples
+  int M;
   ///normalization factor for n(k)
   RealType norm_nofK;
-  ///random generator
-  RandomGenerator_t myRNG;
   ///list of k-points in Cartesian Coordinates
   std::vector<PosType> kPoints;
   ///weight of k-points (make use of symmetry)

From a3ce3f5aaa304239fd4c67cc5ef4cfbe3ece68dd Mon Sep 17 00:00:00 2001
From: Ye Luo <yeluo@anl.gov>
Date: Thu, 23 Sep 2021 18:36:19 -0500
Subject: [PATCH 6/6] More const in MomentumDistribution member.

---
 src/Estimators/MomentumDistribution.cpp | 59 ++++++++++++++-----------
 src/Estimators/MomentumDistribution.h   | 20 ++++++---
 src/Estimators/OperatorEstBase.h        | 12 +++--
 src/Estimators/SpinDensityNew.cpp       | 10 +++--
 src/Estimators/SpinDensityNew.h         |  5 ++-
 5 files changed, 63 insertions(+), 43 deletions(-)

diff --git a/src/Estimators/MomentumDistribution.cpp b/src/Estimators/MomentumDistribution.cpp
index da35a64cb2..166c4df142 100644
--- a/src/Estimators/MomentumDistribution.cpp
+++ b/src/Estimators/MomentumDistribution.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) 2020 QMCPACK developers.
+// Copyright (c) 2021 QMCPACK developers.
 //
 // File developed by: Jaron T. Krogel, krogeljt@ornl.gov, Oak Ridge National Laboratory
 //
@@ -19,17 +19,22 @@
 
 namespace qmcplusplus
 {
-MomentumDistribution::MomentumDistribution(MomentumDistributionInput&& mdi, size_t np, const PosType& twist_in, const LatticeType& lattice, DataLocality dl)
-    : OperatorEstBase(dl), input_(std::move(mdi)), twist(twist_in), Lattice(lattice)
+MomentumDistribution::MomentumDistribution(MomentumDistributionInput&& mdi,
+                                           size_t np,
+                                           const PosType& twist_in,
+                                           const LatticeType& lattice,
+                                           DataLocality dl)
+    : OperatorEstBase(dl),
+      input_(std::move(mdi)),
+      twist(twist_in),
+      Lattice(lattice),
+      M(input_.get<int>("samples")),
+      norm_nofK(1.0 / RealType(M))
 {
-  data_locality_ = dl;
   psi_ratios.resize(np);
 
   myName = input_.get<std::string>("name");
-  M      = input_.get<int>("samples");
 
-  norm_nofK = 1.0 / RealType(M);
-  
   //maximum k-value in the k-grid in cartesian coordinates
   auto kmax = input_.get<RealType>("kmax");
   //maximum k-values in the k-grid along the reciprocal cell axis
@@ -45,7 +50,7 @@ MomentumDistribution::MomentumDistribution(MomentumDistributionInput&& mdi, size
   //length of reciprocal lattice vector
   for (int i = 0; i < OHMMS_DIM; i++)
     vec_length[i] = 2.0 * M_PI * std::sqrt(dot(Lattice.Gv[i], Lattice.Gv[i]));
-  PosType kmaxs = {kmax0,kmax1,kmax2};
+  PosType kmaxs      = {kmax0, kmax1, kmax2};
   RealType sum_kmaxs = kmaxs[0] + kmaxs[1] + kmaxs[2];
   RealType sphere_kmax;
   bool sphere      = kmax > 0.0 ? true : false;
@@ -106,7 +111,7 @@ MomentumDistribution::MomentumDistribution(MomentumDistributionInput&& mdi, size
       }
     }
   }
-  app_log() <<"\n  MomentumDistribution named "<<myName<<"\n";
+  app_log() << "\n  MomentumDistribution named " << myName << "\n";
   if (sphere && !directional)
   {
     app_log() << "    Using all k-space points with (kx^2+ky^2+kz^2)^0.5 < " << sphere_kmax
@@ -167,9 +172,8 @@ MomentumDistribution::MomentumDistribution(MomentumDistributionInput&& mdi, size
   psi_ratios_all.resize(M, psi_ratios.size());
 
   // allocate data storage
-  size_t data_size  = nofK.size();
-  data_             = createLocalData(data_size, data_locality_);
-
+  size_t data_size = nofK.size();
+  data_            = createLocalData(data_size, data_locality_);
 }
 
 
@@ -178,21 +182,21 @@ MomentumDistribution* MomentumDistribution::clone()
   auto* md = new MomentumDistribution(*this);
   if (md->data_locality_ == DataLocality::crowd)
   {
-    app_log()<<"MD::clone dl crowd\n";
+    app_log() << "MD::clone dl crowd\n";
     size_t data_size = data_->size();
     md->data_        = createLocalData(data_size, data_locality_);
   }
   else if (md->data_locality_ == DataLocality::rank)
   {
-    app_log()<<"MD::clone dl rank\n";
+    app_log() << "MD::clone dl rank\n";
     assert(data_locality_ == DataLocality::rank);
-    size_t data_size  = 10; // jtk fix
-    md->data_locality_    = DataLocality::queue;
-    md->data_             = createLocalData(data_size, data_locality_);
+    size_t data_size   = 10; // jtk fix
+    md->data_locality_ = DataLocality::queue;
+    md->data_          = createLocalData(data_size, data_locality_);
   }
   else
-    app_log()<<"MD::clone dl other\n";
-  
+    app_log() << "MD::clone dl other\n";
+
   return md;
 }
 
@@ -229,13 +233,15 @@ void MomentumDistribution::startBlock(int steps)
   //}
   //else
   //  app_log()<<"MD::startBlock dl other\n";
-
 }
 
 /** Gets called every step and writes to thread local data.
  *
  */
-void MomentumDistribution::accumulate(const RefVector<MCPWalker>& walkers, const RefVector<ParticleSet>& psets, const RefVector<TrialWaveFunction>& wfns, RandomGenerator_t& rng)
+void MomentumDistribution::accumulate(const RefVector<MCPWalker>& walkers,
+                                      const RefVector<ParticleSet>& psets,
+                                      const RefVector<TrialWaveFunction>& wfns,
+                                      RandomGenerator_t& rng)
 {
   for (int iw = 0; iw < walkers.size(); ++iw)
   {
@@ -247,7 +253,7 @@ void MomentumDistribution::accumulate(const RefVector<MCPWalker>& walkers, const
     const int np = pset.getTotalNum();
     const int nk = kPoints.size();
 
-    // accumulate weight 
+    // accumulate weight
     //  (required by all estimators, otherwise inf results)
     walkers_weight_ += weight;
 
@@ -263,7 +269,7 @@ void MomentumDistribution::accumulate(const RefVector<MCPWalker>& walkers, const
       psi.evaluateRatiosAlltoOne(pset, psi_ratios);
       for (int i = 0; i < np; ++i)
         psi_ratios_all[s][i] = psi_ratios[i];
-      
+
       for (int ik = 0; ik < nk; ++ik)
         kdotp[ik] = -dot(kPoints[ik], vPos[s]);
       eval_e2iphi(nk, kdotp.data(), phases_vPos[s].data(0), phases_vPos[s].data(1));
@@ -289,16 +295,15 @@ void MomentumDistribution::accumulate(const RefVector<MCPWalker>& walkers, const
 #pragma omp simd aligned(nofK_here, phases_c, phases_s, phases_vPos_c, phases_vPos_s : QMC_SIMD_ALIGNMENT)
         for (int ik = 0; ik < nk; ++ik)
           nofK_here[ik] += (phases_c[ik] * phases_vPos_c[ik] - phases_s[ik] * phases_vPos_s[ik]) * ratio_c -
-            (phases_s[ik] * phases_vPos_c[ik] + phases_c[ik] * phases_vPos_s[ik]) * ratio_s;
+              (phases_s[ik] * phases_vPos_c[ik] + phases_c[ik] * phases_vPos_s[ik]) * ratio_s;
       }
     }
 
     // accumulate data
     for (int ik = 0; ik < nofK.size(); ++ik)
       (*data_)[ik] += weight * nofK[ik] * norm_nofK;
-
   }
-};
+}
 
 
 void MomentumDistribution::collect(const RefVector<OperatorEstBase>& type_erased_operator_estimators)
@@ -326,7 +331,7 @@ void MomentumDistribution::registerOperatorEstimator(hid_t gid)
   h5o->set_dimensions(ng, 0); // JTK: doesn't seem right
   h5o->open(gid);
   h5o->addProperty(const_cast<std::vector<PosType>&>(kPoints), "kpoints");
-  h5o->addProperty(const_cast<std::vector<int>&>(kWeights), "kweights");    
+  h5o->addProperty(const_cast<std::vector<int>&>(kWeights), "kweights");
 }
 
 
diff --git a/src/Estimators/MomentumDistribution.h b/src/Estimators/MomentumDistribution.h
index 33fcab160b..0d220d5b9a 100644
--- a/src/Estimators/MomentumDistribution.h
+++ b/src/Estimators/MomentumDistribution.h
@@ -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) 2020 QMCPACK developers.
+// Copyright (c) 2021 QMCPACK developers.
 //
 // File developed by: Jaron T. Krogel, krogeljt@ornl.gov, Oak Ridge National Laboratory
 //
@@ -36,15 +36,15 @@ class MomentumDistribution : public OperatorEstBase
 
   //data members set only during construction
   ///input values
-  MomentumDistributionInput input_;
+  const MomentumDistributionInput input_;
   ///twist angle
   const PosType twist;
   ///lattice vector
   const LatticeType Lattice;
   ///number of samples
-  int M;
+  const int M;
   ///normalization factor for n(k)
-  RealType norm_nofK;
+  const RealType norm_nofK;
   ///list of k-points in Cartesian Coordinates
   std::vector<PosType> kPoints;
   ///weight of k-points (make use of symmetry)
@@ -69,7 +69,11 @@ class MomentumDistribution : public OperatorEstBase
 
   /** Constructor for MomentumDistributionInput 
    */
-  MomentumDistribution(MomentumDistributionInput&& mdi, size_t np, const PosType& twist, const LatticeType& lattice, DataLocality dl = DataLocality::crowd);
+  MomentumDistribution(MomentumDistributionInput&& mdi,
+                       size_t np,
+                       const PosType& twist,
+                       const LatticeType& lattice,
+                       DataLocality dl = DataLocality::crowd);
 
   //MomentumDistribution(const MomentumDistribution& md);
 
@@ -83,7 +87,10 @@ class MomentumDistribution : public OperatorEstBase
 
   /** accumulate 1 or more walkers of MomentumDistribution samples
    */
-  void accumulate(const RefVector<MCPWalker>& walkers, const RefVector<ParticleSet>& psets, const RefVector<TrialWaveFunction>& wfns, RandomGenerator_t& rng) override;
+  void accumulate(const RefVector<MCPWalker>& walkers,
+                  const RefVector<ParticleSet>& psets,
+                  const RefVector<TrialWaveFunction>& wfns,
+                  RandomGenerator_t& rng) override;
 
   /** this allows the EstimatorManagerNew to reduce without needing to know the details
    *  of MomentumDistribution's data.
@@ -107,7 +114,6 @@ class MomentumDistribution : public OperatorEstBase
    *  big state big coupling design.
    */
   void registerOperatorEstimator(hid_t gid) override;
-
 };
 
 } // namespace qmcplusplus
diff --git a/src/Estimators/OperatorEstBase.h b/src/Estimators/OperatorEstBase.h
index de70e57182..73d17ecd11 100644
--- a/src/Estimators/OperatorEstBase.h
+++ b/src/Estimators/OperatorEstBase.h
@@ -42,7 +42,7 @@ class OperatorEstBase
    */
   using Data = UPtr<std::vector<QMCT::RealType>>;
 
-  /// locality for accumulation data
+  /// locality for accumulation data. FIXME full documentation of this state machine.
   DataLocality data_locality_;
 
   ///name of this object
@@ -63,7 +63,10 @@ class OperatorEstBase
    *  Depending on data locality the accumlation of the result may be different from
    *  the single thread write directly into the OperatorEstimator data.
    */
-  virtual void accumulate(const RefVector<MCPWalker>& walkers, const RefVector<ParticleSet>& psets, const RefVector<TrialWaveFunction>& wfns, RandomGenerator_t& rng) = 0;
+  virtual void accumulate(const RefVector<MCPWalker>& walkers,
+                          const RefVector<ParticleSet>& psets,
+                          const RefVector<TrialWaveFunction>& wfns,
+                          RandomGenerator_t& rng) = 0;
 
   /** Reduce estimator result data from crowds to rank
    *
@@ -78,7 +81,7 @@ class OperatorEstBase
   virtual void normalize(QMCT::RealType invToWgt);
 
   virtual void startBlock(int steps) = 0;
-  
+
   std::vector<QMCT::RealType>& get_data_ref() { return *data_; }
 
   Data& get_data() { return data_; };
@@ -107,10 +110,11 @@ class OperatorEstBase
   /** Return the total walker weight for this block
    */
   QMCT::FullPrecRealType get_walkers_weight() { return walkers_weight_; }
+
 protected:
   QMCT::FullPrecRealType walkers_weight_;
 
-  // convenient Descriptors hdf5 for Operator Estimators only populated for rank scope OperatorEstimator  
+  // convenient Descriptors hdf5 for Operator Estimators only populated for rank scope OperatorEstimator
   UPtrVector<ObservableHelper> h5desc_;
 
   /** create the typed data block for the Operator.
diff --git a/src/Estimators/SpinDensityNew.cpp b/src/Estimators/SpinDensityNew.cpp
index 58c3832feb..9532d10d31 100644
--- a/src/Estimators/SpinDensityNew.cpp
+++ b/src/Estimators/SpinDensityNew.cpp
@@ -21,8 +21,7 @@ namespace qmcplusplus
 SpinDensityNew::SpinDensityNew(SpinDensityInput&& input, const SpeciesSet& species, DataLocality dl)
     : OperatorEstBase(dl), input_(std::move(input)), species_(species), species_size_(getSpeciesSize(species))
 {
-  myName         = "SpinDensity";
-  data_locality_ = dl;
+  myName = "SpinDensity";
 
   if (input_.get_cell().explicitly_defined == true)
     lattice_ = input_.get_cell();
@@ -123,7 +122,10 @@ void SpinDensityNew::startBlock(int steps)
  *  I tried for readable and not doing the optimizers job.
  *  The offsets into bare data are already bad enough.
  */
-void SpinDensityNew::accumulate(const RefVector<MCPWalker>& walkers, const RefVector<ParticleSet>& psets, const RefVector<TrialWaveFunction>& wfns, RandomGenerator_t& rng)
+void SpinDensityNew::accumulate(const RefVector<MCPWalker>& walkers,
+                                const RefVector<ParticleSet>& psets,
+                                const RefVector<TrialWaveFunction>& wfns,
+                                RandomGenerator_t& rng)
 {
   auto& dp_ = derived_parameters_;
   for (int iw = 0; iw < walkers.size(); ++iw)
@@ -147,7 +149,7 @@ void SpinDensityNew::accumulate(const RefVector<MCPWalker>& walkers, const RefVe
         accumulateToData(point, weight);
       }
   }
-};
+}
 
 void SpinDensityNew::accumulateToData(size_t point, QMCT::RealType weight)
 {
diff --git a/src/Estimators/SpinDensityNew.h b/src/Estimators/SpinDensityNew.h
index 32f577dcb1..40f9098693 100644
--- a/src/Estimators/SpinDensityNew.h
+++ b/src/Estimators/SpinDensityNew.h
@@ -66,7 +66,10 @@ class SpinDensityNew : public OperatorEstBase
 
   /** accumulate 1 or more walkers of SpinDensity samples
    */
-  void accumulate(const RefVector<MCPWalker>& walkers, const RefVector<ParticleSet>& psets, const RefVector<TrialWaveFunction>& wfns, RandomGenerator_t& rng) override;
+  void accumulate(const RefVector<MCPWalker>& walkers,
+                  const RefVector<ParticleSet>& psets,
+                  const RefVector<TrialWaveFunction>& wfns,
+                  RandomGenerator_t& rng) override;
 
   /** this allows the EstimatorManagerNew to reduce without needing to know the details
    *  of SpinDensityNew's data.