Batched Structure Factor Estimator

src/Estimators/CMakeLists.txt

@@ -42,7 +42,8 @@ set(QMCEST_SRC
     NEReferencePoints.cpp
     NESpaceGrid.cpp
     EnergyDensityEstimator.cpp
-    StructureFactorInput.cpp)
+    StructureFactorInput.cpp
+    StructureFactorEstimator.cpp)
 
 ####################################
 # create libqmcestimators

src/Estimators/StructureFactorEstimator.cpp

@@ -0,0 +1,114 @@
+//////////////////////////////////////////////////////////////////////////////////////
+// This file is distributed under the University of Illinois/NCSA Open Source License.
+// See LICENSE file in top directory for details.
+//
+// Copyright (c) 2024 QMCPACK developers.
+//
+// File developed by: Peter W. Doak, [email protected], Oak Ridge National Laboratory
+//
+// Some code refactored from: QMCHamiltonian/{SkEstimator.cpp, SkAllEstimator.cpp}
+//////////////////////////////////////////////////////////////////////////////////////
+
+#include "StructureFactorEstimator.h"
+#include "ParticleSet.h"
+#include <LongRange/StructFact.h>
+
+namespace qmcplusplus
+{
+
+StructureFactorEstimator::StructureFactorEstimator(StructureFactorInput& sfi,
+                                                   ParticleSet& pset_ions,
+                                                   ParticleSet& pset_elec,
+                                                   DataLocality data_locality)
+    : OperatorEstBase(data_locality), elns_(pset_elec), ions_(pset_ions)
+{
+  my_name_      = "StructureFactorEstimator";
+  elec_species_ = elns_.getSpeciesSet().getTotalNum();
+  ion_species_  = ions_.getSpeciesSet().getTotalNum();
+
+  num_kpoints_ = pset_ions.getSimulationCell().getKLists().numk;
+  kshell_offsets_   = pset_ions.getSimulationCell().getKLists().kshell;
+  int max_kshell    = kshell_offsets_.size() - 1;
+
+  rhok_tot_r_.resize(num_kpoints_);
+  rhok_tot_i_.resize(num_kpoints_);
+  sfk_e_e_.resize(num_kpoints_);
+  rhok_e_.resize(num_kpoints_);
+  // Legacy comment, but I don't trust it.
+  //for values, we are including e-e structure factor, and e-Ion.  So a total of NumIonSpecies+1 structure factors.
+  //+2 for the real and imaginary parts of rho_k^e
+  //
+  // skAll seems to be written for e-e sf + complex rho_k^e
+  data_.resize(3 * num_kpoints_);
+  kmags_.resize(max_kshell);
+  one_over_degeneracy_kshell_.resize(max_kshell);
+  for (int ks = 0; ks < max_kshell; ks++)
+  {
+    kmags_[ks]                      = std::sqrt(pset_elec.getSimulationCell().getKLists().ksq[kshell_offsets_[ks]]);
+    one_over_degeneracy_kshell_[ks] = 1.0 / static_cast<Real>(kshell_offsets_[ks + 1] - kshell_offsets_[ks]);
+  };
+}
+
+void StructureFactorEstimator::accumulate(const RefVector<MCPWalker>& walkers,
+                                          const RefVector<ParticleSet>& psets,
+                                          const RefVector<TrialWaveFunction>& wfns,
+                                          const RefVector<QMCHamiltonian>& hams,
+                                          RandomBase<FullPrecReal>& rng)
+{
+  for (int iw = 0; iw < walkers.size(); ++iw)
+  {
+    Real weight = walkers[iw].get().Weight;
+
+    //sum over species
+    std::copy(psets[iw].get().getSK().rhok_r[0], psets[iw].get().getSK().rhok_r[0] + num_kpoints_, rhok_tot_r_.begin());
+    std::copy(psets[iw].get().getSK().rhok_i[0], psets[iw].get().getSK().rhok_i[0] + num_kpoints_, rhok_tot_i_.begin());
+    for (int i = 1; i < elec_species_; ++i)
+      accumulate_elements(psets[iw].get().getSK().rhok_r[i], psets[iw].get().getSK().rhok_r[i] + num_kpoints_,
+                          rhok_tot_r_.begin());
+    for (int i = 1; i < elec_species_; ++i)
+      accumulate_elements(psets[iw].get().getSK().rhok_i[i], psets[iw].get().getSK().rhok_i[i] + num_kpoints_,
+                          rhok_tot_i_.begin());
+
+    for (int k = 0; k < num_kpoints_; k++)
+    {
+      sfk_e_e_[k] += weight * (rhok_tot_r_[k] * rhok_tot_r_[k] + rhok_tot_i_[k] * rhok_tot_i_[k]);
+      rhok_e_[k] += weight * std::complex<Real>{rhok_tot_r_[k], rhok_tot_i_[k]};
+    }
+  }
+}
+
+void StructureFactorEstimator::registerOperatorEstimator(hdf_archive& file)
+{
+  hdf_path hdf_name{my_name_};
+  hdf_path path_variables = hdf_name / std::string_view("kpoints");
+  file.push(path_variables, true);
+  file.write(const_cast<std::vector<KPt>&>(ions_.getSimulationCell().getKLists().kpts_cart), "value");
+  file.pop();
+}
+
+void StructureFactorEstimator::write(hdf_archive& file)
+{
+  hdf_path hdf_name{my_name_};
+  file.push(hdf_name);
+  // this is call rhok_e_e in the output of the legacy, but that is just wrong it is |rhok_e_e_|^2
+  file.write(sfk_e_e_, "sfk_e_e");
+  file.write(rhok_e_, "rhok_e_");
+}
+
+void StructureFactorEstimator::collect(const RefVector<OperatorEstBase>& type_erased_operator_estimators)
+{
+  int num_crowds = type_erased_operator_estimators.size();
+  for (OperatorEstBase& crowd_oeb : type_erased_operator_estimators)
+  {
+    StructureFactorEstimator& crowd_sfe = dynamic_cast<StructureFactorEstimator&>(crowd_oeb);
+    this->sfk_e_e_ += crowd_sfe.sfk_e_e_;
+    this->rhok_e_ += crowd_sfe.rhok_e_;
+  }
+  OperatorEstBase::collect(type_erased_operator_estimators);
+}
+
+void StructureFactorEstimator::startBlock(int steps) {}
+
+UPtr<OperatorEstBase> StructureFactorEstimator::spawnCrowdClone() const {}
+
+} // namespace qmcplusplus

src/Estimators/StructureFactorEstimator.h

@@ -0,0 +1,100 @@
+//////////////////////////////////////////////////////////////////////////////////////
+// This file is distributed under the University of Illinois/NCSA Open Source License.
+// See LICENSE file in top directory for details.
+//
+// Copyright (c) 2024 QMCPACK developers.
+//
+// File developed by: Peter W. Doak, [email protected], Oak Ridge National Laboratory
+//
+// Some code refactored from: QMCHamiltonian/{SkEstimator.h, SkAllEstimator.h}
+//////////////////////////////////////////////////////////////////////////////////////
+
+#ifndef QMCPLUSPLUS_STRUCTUREFACTORESTIMATOR_H
+#define QMCPLUSPLUS_STRUCTUREFACTORESTIMATOR_H
+
+#include "OperatorEstBase.h"
+#include "StructureFactorInput.h"
+
+namespace qmcplusplus
+{
+
+namespace testing
+{
+class StructureFactorAccess;
+}
+
+class StructureFactorEstimator : public OperatorEstBase
+{
+public:
+  using QMCT         = QMCTraits;
+  using Real         = QMCT::RealType;
+  using FullPrecReal = QMCT::FullPrecRealType;
+  using KPt          = TinyVector<Real, QMCTraits::DIM>;
+
+  StructureFactorEstimator(StructureFactorInput& sfi,
+                           ParticleSet& pset_ions,
+                           ParticleSet& pset_elec,
+                           DataLocality data_locality = DataLocality::crowd);
+
+  /** accumulate 1 or more walkers of EnergyDensity samples
+   */
+  void accumulate(const RefVector<MCPWalker>& walkers,
+                  const RefVector<ParticleSet>& psets,
+                  const RefVector<TrialWaveFunction>& wfns,
+                  const RefVector<QMCHamiltonian>& hams,
+                  RandomBase<FullPrecReal>& rng) override;
+
+  /** start block entry point
+   */
+  void startBlock(int steps) override;
+
+  UPtr<OperatorEstBase> spawnCrowdClone() const override;
+
+  void registerOperatorEstimator(hdf_archive& file) override;
+  void write(hdf_archive& file);
+  void collect(const RefVector<OperatorEstBase>& type_erased_operator_estimators) override;
+
+  long long getNumKPoints() { return num_kpoints_; }
+protected:
+  // Testing functions
+  const Vector<Real>& getSKElecElec() const { return sfk_e_e_; }
+  const Vector<std::complex<Real>>& getRhoKElec() const { return rhok_e_; }
+
+private:
+  int elec_species_;
+  ParticleSet& elns_;
+  int ion_species_;
+  ParticleSet& ions_;
+
+  /// number of k points
+  long long num_kpoints_;
+
+  // std::vector<int> kshell_degeneracy_;
+  /// kpts which belong to the ith-shell [kshell[i], kshell[i+1])
+  std::vector<int> kshell_offsets_;
+
+  /// All the following are indexed by kshell
+  /// instantaneous structure factor for a k shell
+  std::vector<Real> kmags_;
+
+  std::vector<Real> one_over_degeneracy_kshell_;
+
+  /// Accumulated, its clearer to do it this way that use the base class data_ but means we don't use that base class infrastructure
+  Vector<Real> sfk_e_e_;
+  Vector<std::complex<Real>> rhok_e_;
+
+  /*@{
+   *  work area to reduce over electron species, they should probably just be complex as well.
+   *  but particle set stores the real and imaginary parts as two real vectors.
+   */
+  Vector<Real> rhok_tot_r_;
+  Vector<Real> rhok_tot_i_;
+  /*@}*/
+
+public:
+  friend class testing::StructureFactorAccess;
+};
+
+} // namespace qmcplusplus
+
+#endif

src/Estimators/tests/CMakeLists.txt

@@ -70,6 +70,7 @@ set(SRCS
     test_EnergyDensityInput.cpp
     test_EnergyDensityEstimator.cpp
     test_StructureFactorInput.cpp
+    test_StructureFactorEstimator.cpp
     )
 
 add_executable(${UTEST_EXE} ${SRCS})

src/Estimators/tests/GenerateRandomParticleSets.cpp

@@ -32,7 +32,7 @@ std::vector<ParticleSet> generateRandomParticleSets(ParticleSet& pset_target,
   std::vector<ParticleSet> psets(num_psets, pset_target);
   if (generate_test_data)
   {
-    std::cout << "Initialize OneBodyDensityMatrices::accumulate psets with:\n{";
+    std::cout << "Initialize psets with:\n{";
     std::vector<ParticleSet> psets;
     for (int iw = 0; iw < nwalkers; ++iw)
     {