Merge pull request #5389 from PDoakORNL/kcontainer_working_precision

kcontainer now encapsulated and provides k elements at working prec

Author: prckent

src/Particle/LongRange/EwaldHandler2D.cpp

@@ -34,16 +34,18 @@ void EwaldHandler2D::fillFk(const KContainer& KList)
   const mRealType kalpha = 1.0 / (2.0*alpha);
   mRealType kmag, uk;
 
-  Fk.resize(KList.kpts_cart.size());
-  MaxKshell = KList.kshell.size() - 1;
+  Fk.resize(KList.getKptsCartWorking().size());
+  const auto& kshell = KList.getKShell();
+  MaxKshell = kshell.size() - 1;
   Fk_symm.resize(MaxKshell);
 
+  const auto& ksq = KList.getKSQWorking();
   for (int ks = 0, ki = 0; ks < Fk_symm.size(); ks++)
   {
-    kmag = std::sqrt(KList.ksq[ki]);
+    kmag = std::sqrt(ksq[ki]);
     uk = knorm * erfc(kalpha*kmag)/kmag;
     Fk_symm[ks] = uk;
-    while (ki < KList.kshell[ks + 1] && ki < Fk.size())
+    while (ki < kshell[ks + 1] && ki < Fk.size())
       Fk[ki++] = uk;
   }
 }

src/Particle/LongRange/EwaldHandler3D.cpp

@@ -53,22 +53,24 @@ EwaldHandler3D::EwaldHandler3D(const EwaldHandler3D& aLR, ParticleSet& ref)
 
 void EwaldHandler3D::fillFk(const KContainer& KList)
 {
-  Fk.resize(KList.kpts_cart.size());
-  Fkg.resize(KList.kpts_cart.size());
-  const std::vector<int>& kshell(KList.kshell);
+  const auto& kpts_cart = KList.getKptsCartWorking();
+  Fk.resize(kpts_cart.size());
+  Fkg.resize(kpts_cart.size());
+  const std::vector<int>& kshell(KList.getKShell());
 
   MaxKshell = kshell.size() - 1;
 
   Fk_symm.resize(MaxKshell);
   kMag.resize(MaxKshell);
   mRealType kgauss = 1.0 / (4 * Sigma * Sigma);
   mRealType knorm  = 4 * M_PI / Volume;
+  const auto ksq = KList.getKSQWorking();
   for (int ks = 0, ki = 0; ks < Fk_symm.size(); ks++)
   {
-    mRealType t2e = KList.ksq[ki] * kgauss;
-    mRealType uk  = knorm * std::exp(-t2e) / KList.ksq[ki];
+    mRealType t2e = ksq[ki] * kgauss;
+    mRealType uk  = knorm * std::exp(-t2e) / ksq[ki];
     Fk_symm[ks]   = uk;
-    while (ki < KList.kshell[ks + 1] && ki < Fk.size())
+    while (ki < kshell[ks + 1] && ki < Fk.size())
       Fk[ki++] = uk;
   }
 

src/Particle/LongRange/EwaldHandler3D.h

@@ -103,24 +103,26 @@ class EwaldHandler3D : public LRHandlerBase
 
   void fillYkgstrain(const KContainer& KList)
   {
-    Fkgstrain.resize(KList.kpts_cart.size());
-    const std::vector<int>& kshell(KList.kshell);
+    Fkgstrain.resize(KList.getKptsCartWorking().size());
+    const std::vector<int>& kshell(KList.getKShell());
     MaxKshell = kshell.size() - 1;
+    const auto& ksq = KList.getKSQWorking();
     for (int ks = 0, ki = 0; ks < MaxKshell; ks++)
     {
-      mRealType uk = evalYkgstrain(std::sqrt(KList.ksq[ki]));
-      while (ki < KList.kshell[ks + 1] && ki < Fkgstrain.size())
+      mRealType uk = evalYkgstrain(std::sqrt(ksq[ki]));
+      while (ki < kshell[ks + 1] && ki < Fkgstrain.size())
         Fkgstrain[ki++] = uk;
     }
   }
 
   void filldFk_dk(const KContainer& KList)
   {
-    dFk_dstrain.resize(KList.kpts_cart.size());
-
+    const auto& kpts_cart = KList.getKptsCartWorking();
+    dFk_dstrain.resize(kpts_cart.size());
+    const auto& ksq = KList.getKSQWorking();
     for (int ki = 0; ki < dFk_dstrain.size(); ki++)
     {
-      dFk_dstrain[ki] = evaluateLR_dstrain(KList.kpts_cart[ki], std::sqrt(KList.ksq[ki]));
+      dFk_dstrain[ki] = evaluateLR_dstrain(kpts_cart[ki], std::sqrt(ksq[ki]));
     }
   }
 

src/Particle/LongRange/EwaldHandlerQuasi2D.cpp

@@ -35,19 +35,19 @@ void EwaldHandlerQuasi2D::fillFk(const KContainer& KList)
 {
   const mRealType knorm = M_PI / area;
   mRealType kmag, uk;
-
-  Fk.resize(KList.kpts_cart.size());
-  MaxKshell = KList.kshell.size() - 1;
+  const auto& kpts_cart = KList.getKptsCartWorking();
+  Fk.resize(kpts_cart.size());
+  MaxKshell = KList.getKShell().size() - 1;
   Fk_symm.resize(MaxKshell);
 
   kmags.resize(MaxKshell);
   for (int ks = 0, ki = 0; ks < Fk_symm.size(); ks++)
   {
-    kmag = std::sqrt(KList.ksq[ki]);
+    kmag = std::sqrt(KList.getKSQWorking()[ki]);
     kmags[ks] = kmag; // store k magnitutes
     uk = knorm/kmag;
     Fk_symm[ks] = uk;
-    while (ki < KList.kshell[ks + 1] && ki < Fk.size())
+    while (ki < KList.getKShell()[ks + 1] && ki < Fk.size())
       Fk[ki++] = uk;
   }
 }

src/Particle/LongRange/KContainer.cpp

@@ -2,11 +2,12 @@
 // This file is distributed under the University of Illinois/NCSA Open Source License.
 // See LICENSE file in top directory for details.
 //
-// Copyright (c) 2016 Jeongnim Kim and QMCPACK developers.
+// Copyright (c) 2025 QMCPACK developers.
 //
 // File developed by: Jeremy McMinnis, jmcminis@gmail.com, University of Illinois at Urbana-Champaign
 //                    Jeongnim Kim, jeongnim.kim@gmail.com, University of Illinois at Urbana-Champaign
 //                    Mark A. Berrill, berrillma@ornl.gov, Oak Ridge National Laboratory
+//                    Peter W. Doak, doakpw@ornl.gov, Oak Ridge National Laboratory
 //
 // File created by: Jeongnim Kim, jeongnim.kim@gmail.com, University of Illinois at Urbana-Champaign
 //////////////////////////////////////////////////////////////////////////////////////
@@ -21,11 +22,41 @@
 
 namespace qmcplusplus
 {
-void KContainer::updateKLists(const ParticleLayout& lattice,
-                              RealType kc,
-                              unsigned ndim,
-                              const PosType& twist,
-                              bool useSphere)
+
+template<typename REAL>
+const std::vector<typename KContainerT<REAL>::AppPosition>& KContainerT<
+    REAL>::getKptsCartWorking() const
+{
+  // This is an `if constexpr` so it should not cost a branch at runtime.
+  if constexpr (std::is_same_v<decltype(kpts_cart_), decltype(kpts_cart_working_)>)
+    return kpts_cart_;
+  else
+    return kpts_cart_working_;
+}
+
+template<typename REAL>
+const std::vector<REAL>& KContainerT<REAL>::getKSQWorking() const
+{
+  // This is an `if constexpr` so it should not cost a branch at runtime.
+  if constexpr (std::is_same<decltype(ksq_), decltype(ksq_working_)>::value)
+    return ksq_;
+  else
+    return ksq_working_;
+}
+
+template<typename REAL>
+int KContainerT<REAL>::getMinusK(int k) const
+{
+  assert(k < minusk.size());
+  return minusk[k];
+}
+
+template<typename REAL>
+void KContainerT<REAL>::updateKLists(const Lattice& lattice,
+                                     FullPrecReal kc,
+                                     unsigned ndim,
+                                     const Position& twist,
+                                     bool useSphere)
 {
   kcutoff = kc;
   if (kcutoff <= 0.0)
@@ -37,11 +68,12 @@ void KContainer::updateKLists(const ParticleLayout& lattice,
 
   app_log() << "  KContainer initialised with cutoff " << kcutoff << std::endl;
   app_log() << "   # of K-shell  = " << kshell.size() << std::endl;
-  app_log() << "   # of K points = " << kpts.size() << std::endl;
+  app_log() << "   # of K points = " << kpts_.size() << std::endl;
   app_log() << std::endl;
 }
 
-void KContainer::findApproxMMax(const ParticleLayout& lattice, unsigned ndim)
+template<typename REAL>
+void KContainerT<REAL>::findApproxMMax(const Lattice& lattice, unsigned ndim)
 {
   //Estimate the size of the parallelpiped that encompasses a sphere of kcutoff.
   //mmax is stored as integer translations of the reciprocal cell vectors.
@@ -102,19 +134,22 @@ void KContainer::findApproxMMax(const ParticleLayout& lattice, unsigned ndim)
     mmax[1] = 0;
 }
 
-void KContainer::BuildKLists(const ParticleLayout& lattice, const PosType& twist, bool useSphere)
+template<typename REAL>
+void KContainerT<REAL>::BuildKLists(const Lattice& lattice,
+                                    const Position& twist,
+                                    bool useSphere)
 {
   TinyVector<int, DIM + 1> TempActualMax;
   TinyVector<int, DIM> kvec;
-  TinyVector<RealType, DIM> kvec_cart;
-  RealType modk2;
+  TinyVector<FullPrecReal, DIM> kvec_cart;
+  FullPrecReal modk2;
   std::vector<TinyVector<int, DIM>> kpts_tmp;
-  std::vector<PosType> kpts_cart_tmp;
-  std::vector<RealType> ksq_tmp;
+  std::vector<PositionFull> kpts_cart_tmp;
+  std::vector<FullPrecReal> ksq_tmp;
   // reserve the space for memory efficiency
   if (useSphere)
   {
-    const RealType kcut2 = kcutoff * kcutoff;
+    const FullPrecReal kcut2 = kcutoff * kcutoff;
     //Loop over guesses for valid k-points.
     for (int i = -mmax[0]; i <= mmax[0]; i++)
     {
@@ -208,10 +243,10 @@ void KContainer::BuildKLists(const ParticleLayout& lattice, const PosType& twist
     }
   }
   std::map<int64_t, std::vector<int>*>::iterator it(kpts_sorted.begin());
-  kpts.resize(numk);
-  kpts_cart.resize(numk);
+  kpts_.resize(numk);
+  kpts_cart_.resize(numk);
   kpts_cart_soa_.resize(numk);
-  ksq.resize(numk);
+  ksq_.resize(numk);
   kshell.resize(kpts_sorted.size() + 1, 0);
   int ok = 0, ish = 0;
   while (it != kpts_sorted.end())
@@ -220,10 +255,10 @@ void KContainer::BuildKLists(const ParticleLayout& lattice, const PosType& twist
     while (vit != (*it).second->end())
     {
       int ik             = (*vit);
-      kpts[ok]           = kpts_tmp[ik];
-      kpts_cart[ok]      = kpts_cart_tmp[ik];
+      kpts_[ok]          = kpts_tmp[ik];
+      kpts_cart_[ok]     = kpts_cart_tmp[ik];
       kpts_cart_soa_(ok) = kpts_cart_tmp[ik];
-      ksq[ok]            = ksq_tmp[ik];
+      ksq_[ok]           = ksq_tmp[ik];
       ++vit;
       ++ok;
     }
@@ -232,6 +267,13 @@ void KContainer::BuildKLists(const ParticleLayout& lattice, const PosType& twist
     ++ish;
   }
   kpts_cart_soa_.updateTo();
+  if constexpr (!std::is_same<Real, FullPrecReal>::value)
+  {
+    // This copy implicity does the precision reduction.
+    // the working vectors are not used or initialized for full precision builds.
+    std::copy(kpts_cart_.begin(), kpts_cart_.end(), std::back_inserter(kpts_cart_working_));
+    std::copy(ksq_.begin(), ksq_.end(), std::back_inserter(ksq_working_));
+  }
   it = kpts_sorted.begin();
   std::map<int64_t, std::vector<int>*>::iterator e_it(kpts_sorted.end());
   while (it != e_it)
@@ -262,13 +304,19 @@ void KContainer::BuildKLists(const ParticleLayout& lattice, const PosType& twist
   std::map<int64_t, int> hashToIndex;
   for (int ki = 0; ki < numk; ki++)
   {
-    hashToIndex[getHashOfVec(kpts[ki], numk)] = ki;
+    hashToIndex[getHashOfVec(kpts_[ki], numk)] = ki;
   }
   // Use the map to find the index of -k from the index of k
   for (int ki = 0; ki < numk; ki++)
   {
-    minusk[ki] = hashToIndex[getHashOfVec(-1 * kpts[ki], numk)];
+    minusk[ki] = hashToIndex[getHashOfVec(-1 * kpts_[ki], numk)];
   }
 }
 
+#ifdef MIXED_PRECISION
+template class KContainerT<float>;
+template class KContainerT<double>;
+#else
+template class KContainerT<double>;
+#endif
 } // namespace qmcplusplus