Merge pull request QMCPACK#5848 from ye-luo/refactor-splines

Major refactoring of spline SPO related classes

Author: prckent

src/QMCWaveFunctions/BandInfo.cpp

@@ -15,6 +15,8 @@
  */
 #include "BandInfo.h"
 #include "QMCWaveFunctions/SPOSetInfo.h"
+#include <Message/UniformCommunicateError.h>
+
 namespace qmcplusplus
 {
 BandInfoGroup::BandInfoGroup() : NumSPOs(0), FirstBand(0) {}
@@ -24,21 +26,11 @@ void BandInfoGroup::selectBands(const std::vector<BandInfo>& bigspace, int first
   app_log() << "BandInfoGroup::selectBands bigspace has " << bigspace.size() << " distinct orbitals " << std::endl;
   myBands.clear();
 
-  int iorb    = 0;
-  int N       = bigspace.size();
-  int n_lower = 0;
-  do
-  {
-    if (iorb >= first_orb)
-      break;
-    n_lower += (bigspace[iorb].MakeTwoCopies) ? 2 : 1;
-    ++iorb;
-  } while (iorb < N);
+  int iorb    = first_orb;
+  const int N = bigspace.size();
 
   if (iorb >= N)
-  {
-    APP_ABORT("BandInfoGroup::selectBands failed due to iorb>=N");
-  }
+    throw UniformCommunicateError("BandInfoGroup::selectBands failed due to iorb>=N");
 
   FirstBand  = iorb;
   NumSPOs    = 0;
@@ -55,6 +47,9 @@ void BandInfoGroup::selectBands(const std::vector<BandInfo>& bigspace, int first
   app_log() << "  Number of distinct bands " << myBands.size() << std::endl;
   app_log() << "  First Band index " << FirstBand << std::endl;
   app_log() << "  Size of SPOs " << NumSPOs << std::endl;
+
+  if (NumSPOs != num_spos)
+    throw UniformCommunicateError("Insufficient bands to generate SPOs. Requested amount " + std::to_string(num_spos));
 }
 
 } // namespace qmcplusplus

src/QMCWaveFunctions/BandInfo.h

@@ -71,8 +71,6 @@ struct BandInfoGroup
   int NumSPOs;
   ///starting band
   int FirstBand;
-  ///twist index set by the full band not by the subset
-  int TwistIndex;
   ///Bands that belong to this group
   std::vector<BandInfo> myBands;
   ///name of this band

src/QMCWaveFunctions/BsplineFactory/BsplineReader.cpp

@@ -22,14 +22,15 @@
 #include "OhmmsData/AttributeSet.h"
 #include "Message/CommOperators.h"
 #include <PlatformSelector.hpp>
+#include "einspline_helper.hpp"
 
 #include <array>
 #include <filesystem>
 
 namespace qmcplusplus
 {
-BsplineReader::BsplineReader(EinsplineSetBuilder* e)
-    : mybuilder(e), checkNorm(true), saveSplineCoefs(false), rotate(true)
+BsplineReader::BsplineReader(EinsplineSetBuilder* e, bool use_duplex_splines)
+    : mybuilder(e), checkNorm(true), saveSplineCoefs(false), rotate(true), use_duplex_splines_(use_duplex_splines)
 { myComm = mybuilder->getCommunicator(); }
 
 BsplineReader::~BsplineReader() = default;
@@ -107,14 +108,61 @@ std::unique_ptr<SPOSet> BsplineReader::create_spline_set(const std::string& spo_
   }
 
   BandInfoGroup vals;
-  vals.TwistIndex = fullband[0].TwistIndex;
-  vals.GroupID    = 0;
-  vals.myName     = make_bandgroup_name(spo_name, spin, mybuilder->twist_num_, mybuilder->TileMatrix, 0, size);
+  vals.GroupID = 0;
+  vals.myName  = make_bandgroup_name(spo_name, spin, mybuilder->twist_num_, mybuilder->TileMatrix, 0, size);
   vals.selectBands(fullband, 0, size);
 
   return create_spline_set(spo_name, spin, vals);
 }
 
+bool BsplineReader::lookforSplineDataDumpFile(const BandInfoGroup& bandgroup,
+                                              const std::string& keyword,
+                                              size_t datatype_size) const
+{
+  int foundspline = 0;
+  if (myComm->rank() == 0)
+  {
+    hdf_archive h5f(myComm);
+    foundspline = h5f.open(getSplineDumpFileName(bandgroup), H5F_ACC_RDONLY);
+    if (foundspline)
+    {
+      std::string aname("none");
+      foundspline = h5f.readEntry(aname, "class_name");
+      foundspline = (aname.find(keyword) != std::string::npos);
+    }
+    if (foundspline)
+    {
+      int sizeD   = 0;
+      foundspline = h5f.readEntry(sizeD, "sizeof");
+      foundspline = (sizeD == datatype_size);
+    }
+    h5f.close();
+  }
+  myComm->bcast(foundspline);
+  return foundspline;
+}
+
+void BsplineReader::readOneOrbitalCoefs(const std::string& s, hdf_archive& h5f, Vector<std::complex<double>>& cG) const
+{
+  if (!h5f.readEntry(cG, s))
+  {
+    std::ostringstream msg;
+    msg << "SplineSetReader Failed to read band(s) from h5 file. " << "Attempted dataset " << s << " with " << cG.size()
+        << " complex numbers." << std::endl;
+    throw std::runtime_error(msg.str());
+  }
+  double total_norm = compute_norm(cG);
+  if ((checkNorm) && (std::abs(total_norm - 1.0) > PW_COEFF_NORM_TOLERANCE))
+  {
+    std::ostringstream msg;
+    msg << "SplineSetReader The orbital dataset " << s << " has a wrong norm " << total_norm
+        << ", computed from plane wave coefficients!" << std::endl
+        << "This may indicate a problem with the HDF5 library versions used "
+        << "during wavefunction conversion or read." << std::endl;
+    throw std::runtime_error(msg.str());
+  }
+}
+
 std::unique_ptr<SPOSet> BsplineReader::create_spline_set(const std::string& spo_name,
                                                          int spin,
                                                          SPOSetInputInfo& input_info)
@@ -134,10 +182,9 @@ std::unique_ptr<SPOSet> BsplineReader::create_spline_set(const std::string& spo_
   }
 
   BandInfoGroup vals;
-  vals.TwistIndex = fullband[0].TwistIndex;
-  vals.GroupID    = 0;
-  vals.myName     = make_bandgroup_name(spo_name, spin, mybuilder->twist_num_, mybuilder->TileMatrix,
-                                        input_info.min_index(), input_info.max_index());
+  vals.GroupID = 0;
+  vals.myName  = make_bandgroup_name(spo_name, spin, mybuilder->twist_num_, mybuilder->TileMatrix,
+                                     input_info.min_index(), input_info.max_index());
   vals.selectBands(fullband, spo2band[spin][input_info.min_index()], input_info.max_index() - input_info.min_index());
 
   return create_spline_set(spo_name, spin, vals);

src/QMCWaveFunctions/BsplineFactory/BsplineReader.h

@@ -52,7 +52,7 @@ class BsplineReader
                            std::vector<int>& band2spo);
 
 public:
-  BsplineReader(EinsplineSetBuilder* e);
+  BsplineReader(EinsplineSetBuilder* e, bool use_duplex_splines);
 
   virtual ~BsplineReader();
 
@@ -86,6 +86,8 @@ class BsplineReader
   bool rotate;
   ///map from spo index to band index
   std::vector<std::vector<int>> spo2band;
+  /// if true, use two real-valued splines for one complex-valued DFT orbital.
+  const bool use_duplex_splines_;
   /// if true, use offload
   bool use_offload;
 
@@ -100,13 +102,13 @@ class BsplineReader
   /** read gvectors and set the mesh, and prepare for einspline
    */
   template<typename GT, typename BCT>
-  inline void set_grid(const TinyVector<int, 3>& halfg, GT* xyz_grid, BCT* xyz_bc) const
+  static void set_grid(const TinyVector<int, 3>& mesh_sizes, const TinyVector<int, 3>& halfg, GT* xyz_grid, BCT* xyz_bc)
   {
     for (int j = 0; j < 3; ++j)
     {
       xyz_grid[j].start = 0.0;
       xyz_grid[j].end   = 1.0;
-      xyz_grid[j].num   = mybuilder->MeshSize[j];
+      xyz_grid[j].num   = mesh_sizes[j];
 
       if (halfg[j])
       {
@@ -131,8 +133,6 @@ class BsplineReader
     const int N       = bandgroup.getNumDistinctOrbitals();
     const int numOrbs = bandgroup.getNumSPOs();
 
-    bspline.resizeStorage(N);
-
     const std::vector<BandInfo>& cur_bands = bandgroup.myBands;
     for (int iorb = 0, num = 0; iorb < N; iorb++)
     {
@@ -142,23 +142,27 @@ class BsplineReader
       num += bspline.MakeTwoCopies[iorb] ? 2 : 1;
     }
 
+    bspline.resize_kpoints();
     app_log() << "NumDistinctOrbitals " << N << " numOrbs = " << numOrbs << std::endl;
+  }
 
-    bspline.HalfG             = 0;
-    TinyVector<int, 3> bconds = mybuilder->TargetPtcl.getLattice().BoxBConds;
-    if (!bspline.isComplex())
-    {
-      //no k-point folding, single special k point (G, L ...)
-      TinyVector<double, 3> twist0 = mybuilder->primcell_kpoints[bandgroup.TwistIndex];
-      for (int i = 0; i < 3; i++)
-        if (bconds[i] && ((std::abs(std::abs(twist0[i]) - 0.5) < 1.0e-8)))
-          bspline.HalfG[i] = 1;
-        else
-          bspline.HalfG[i] = 0;
-      app_log() << "  TwistIndex = " << cur_bands[0].TwistIndex << " TwistAngle " << twist0 << std::endl;
-      app_log() << "   HalfG = " << bspline.HalfG << std::endl;
-    }
-    app_log().flush();
+  /** compute sign bits at the G/2 boundaries
+   * no supercell, no k-point folding, single special k point (G, L ...)
+   */
+  static TinyVector<int, 3> computeHalfG(const TinyVector<int, OHMMS_DIM>& bconds,
+                                         const std::vector<TinyVector<double, OHMMS_DIM>>& primcell_kpoints,
+                                         size_t twist0_index)
+  {
+    TinyVector<int, 3> halfG;
+    const auto& twist0 = primcell_kpoints[twist0_index];
+    app_log() << "  TwistIndex = " << twist0_index << " TwistAngle " << twist0 << std::endl;
+    for (int i = 0; i < 3; i++)
+      if (bconds[i] && ((std::abs(std::abs(twist0[i]) - 0.5) < 1.0e-8)))
+        halfG[i] = 1;
+      else
+        halfG[i] = 0;
+    app_log() << "   HalfG = " << halfG << std::endl;
+    return halfG;
   }
 
   /** return the path name in hdf5
@@ -172,6 +176,22 @@ class BsplineReader
     path << "/electrons/kpoint_" << ti << "/spin_" << spin << "/state_" << ib << "/psi_g";
     return path.str();
   }
+
+  /** create data space in the spline object and try open spline dump files.
+   * @param bandgroup band info
+   * @param bspline the spline object being worked on
+   * @return true if dumpfile pass class name and data type size check
+   */
+  bool lookforSplineDataDumpFile(const BandInfoGroup& bandgroup,
+                                 const std::string& keyword,
+                                 size_t datatype_size) const;
+
+  /** read planewave coefficients from h5 file
+   * @param s data set full path in h5
+   * @param h5f hdf5 file handle
+   * @param cG vector to store coefficients
+   */
+  void readOneOrbitalCoefs(const std::string& s, hdf_archive& h5f, Vector<std::complex<double>>& cG) const;
 };
 
 } // namespace qmcplusplus

src/QMCWaveFunctions/BsplineFactory/BsplineSet.h

@@ -25,6 +25,12 @@
 
 namespace qmcplusplus
 {
+
+template<class ST>
+class SplineSetReader;
+template<class ST>
+class HybridRepSetReader;
+
 /** BsplineSet is the base class for SplineC2C, SplineC2R, SplineR2R.
  * Its derived template classes manage the storage and evaluation at given precision.
  * BsplineSet also implements a few fallback routines in case optimized implementation is not necessary in the derived class.
@@ -52,7 +58,6 @@ class BsplineSet : public SPOSet
       : SPOSet(my_name, size), prim_lattice_(prim_lattice)
   {}
 
-  virtual bool isComplex() const         = 0;
   virtual std::string getKeyword() const = 0;
 
   auto& getHalfG() const { return HalfG; }
@@ -69,6 +74,9 @@ class BsplineSet : public SPOSet
   /// resize vectors related to spline evaluaton results.
   virtual void resizeStorage(size_t n) = 0;
 
+  /** remap kPoints to pack the double copy */
+  virtual void resize_kpoints() {}
+
   ///remap kpoints to group general kpoints & special kpoints
   int remap_kpoints()
   {
@@ -236,8 +244,10 @@ class BsplineSet : public SPOSet
     //Do nothing, since Einsplines don't explicitly depend on ion positions.
   }
 
-  template<class BSPLINESPO>
+  template<class ST>
   friend class SplineSetReader;
+  template<class ST>
+  friend class HybridRepSetReader;
   friend class BsplineReader;
 };
 

src/QMCWaveFunctions/BsplineFactory/EinsplineSetBuilder_createSPOs.cpp

@@ -29,7 +29,8 @@
 #include "einspline_helper.hpp"
 #include "BsplineReader.h"
 #include "BsplineSet.h"
-#include "createBsplineReader.h"
+#include "SplineSetReader.h"
+#include "HybridRepSetReader.h"
 
 #include <array>
 #include <string_view>
@@ -115,12 +116,8 @@ std::unique_ptr<SPOSet> EinsplineSetBuilder::createSPOSetFromXML(xmlNodePtr cur)
   std::string sourceName;
   std::string spo_prec("double");
   std::string truncate("no");
-  std::string hybrid_rep("no");
+  std::string hybrid_rep;
   std::string skip_checks("no");
-  std::string use_einspline_set_extended(
-      "no"); // use old spline library for high-order derivatives, e.g. needed for backflow optimization
-  std::string useGPU;
-  std::string GPUsharing = "no";
 
   ScopedTimer spo_timer_scope(createGlobalTimer("einspline::CreateSPOSetFromXML", timer_level_medium));
 
@@ -136,9 +133,7 @@ std::unique_ptr<SPOSet> EinsplineSetBuilder::createSPOSetFromXML(xmlNodePtr cur)
     a_root.add(twist_inp, "twist");
     a_root.add(sourceName, "source");
     a_root.add(MeshFactor, "meshfactor");
-    a_root.add(hybrid_rep, "hybridrep");
-    a_root.add(useGPU, "gpu", CPUOMPTargetSelector::candidate_values);
-    a_root.add(GPUsharing, "gpusharing"); // split spline across GPUs visible per rank
+    a_root.add(hybrid_rep, "hybridrep", {"no", "yes"});
     a_root.add(spo_prec, "precision");
     a_root.add(truncate, "truncate");
     a_root.add(skip_checks, "skip_checks");
@@ -244,24 +239,30 @@ std::unique_ptr<SPOSet> EinsplineSetBuilder::createSPOSetFromXML(xmlNodePtr cur)
   if (!ReadGvectors_ESHDF())
     myComm->barrier_and_abort("Failed to load g-vectors.");
 
-  bool use_single = (spo_prec == "single" || spo_prec == "float");
+  const bool use_single = (spo_prec == "single" || spo_prec == "float");
+  app_summary() << "    Using " << (use_single ? "single" : "double") << " precision B-spline coefficients."
+                << std::endl;
+
+  const bool use_hybridrep = hybrid_rep == "yes";
+  app_summary() << "    Using " << (use_hybridrep ? "hybrid" : "regular 3D cublic B-spline")
+                << " orbital representation." << std::endl;
 
   // safeguard for a removed feature
   if (truncate == "yes")
     myComm->barrier_and_abort(
         "The 'truncate' feature of spline SPO has been removed. Please use hybrid orbital representation.");
 
-#if !defined(QMC_COMPLEX)
-  if (use_real_splines_)
-  {
-    if (MixedSplineReader == 0)
-      MixedSplineReader = createBsplineReal(this, use_single, hybrid_rep == "yes");
-  }
-  else
-#endif
+  if (!MixedSplineReader)
   {
-    if (MixedSplineReader == 0)
-      MixedSplineReader = createBsplineComplex(this, use_single, hybrid_rep == "yes", useGPU);
+    if (use_hybridrep)
+      if (use_single)
+        MixedSplineReader = std::make_unique<HybridRepSetReader<float>>(this, !use_real_splines_);
+      else
+        MixedSplineReader = std::make_unique<HybridRepSetReader<double>>(this, !use_real_splines_);
+    else if (use_single)
+      MixedSplineReader = std::make_unique<SplineSetReader<float>>(this, !use_real_splines_);
+    else
+      MixedSplineReader = std::make_unique<SplineSetReader<double>>(this, !use_real_splines_);
   }
 
   MixedSplineReader->setCommon(XMLRoot);