One Body Spin dependent Jastrow

src/QMCWaveFunctions/Jastrow/RadialJastrowBuilder.cpp

@@ -11,8 +11,9 @@
 //////////////////////////////////////////////////////////////////////////////////////
 
 #include "RadialJastrowBuilder.h"
-
+#include <type_traits>
 #include "QMCWaveFunctions/Jastrow/J1OrbitalSoA.h"
+#include "QMCWaveFunctions/Jastrow/J1Spin.h"
 #include "QMCWaveFunctions/Jastrow/J2OrbitalSoA.h"
 
 #if defined(ENABLE_OFFLOAD)
@@ -41,19 +42,19 @@
 
 namespace qmcplusplus
 {
-
 // quick helper class to allow use of RPA
 class RPAFunctor
 {};
 
 // helper class to simplify and localize ugly ifdef stuff for types
-template<class RadFuncType, unsigned Implementation>
+template<class RadFuncType, unsigned Implementation = RadialJastrowBuilder::detail::CPU>
 class JastrowTypeHelper
 {
 public:
-  using J1OrbitalType     = J1OrbitalSoA<RadFuncType>;
-  using J2OrbitalType     = J2OrbitalSoA<RadFuncType>;
-  using DiffJ2OrbitalType = DiffTwoBodyJastrowOrbital<RadFuncType>;
+  using J1Type     = J1OrbitalSoA<RadFuncType>;
+  using J1SpinType = J1Spin<RadFuncType>;
+  using J2Type     = J2OrbitalSoA<RadFuncType>;
+  using DiffJ2Type = DiffTwoBodyJastrowOrbital<RadFuncType>;
 };
 
 #if defined(QMC_CUDA)
@@ -62,30 +63,21 @@ class JastrowTypeHelper<BsplineFunctor<RadialJastrowBuilder::RealType>, RadialJa
 {
 public:
   using RadFuncType = BsplineFunctor<RadialJastrowBuilder::RealType>;
-  using J1OrbitalType = OneBodyJastrowOrbitalBspline<RadFuncType>;
-  using J2OrbitalType = TwoBodyJastrowOrbitalBspline<RadFuncType>;
-  using DiffJ2OrbitalType = DiffTwoBodyJastrowOrbital<RadFuncType>;
+  using J1Type      = OneBodyJastrowOrbitalBspline<RadFuncType>;
+  using J1SpinType  = void;
+  using J2Type      = TwoBodyJastrowOrbitalBspline<RadFuncType>;
+  using DiffJ2Type  = DiffTwoBodyJastrowOrbital<RadFuncType>;
 };
 #endif
 
-template<>
-class JastrowTypeHelper<BsplineFunctor<RadialJastrowBuilder::RealType>, RadialJastrowBuilder::detail::CPU>
-{
-public:
-  using RadFuncType = BsplineFunctor<RadialJastrowBuilder::RealType>;
-  using J1OrbitalType = J1OrbitalSoA<RadFuncType>;
-  using J2OrbitalType = J2OrbitalSoA<RadFuncType>;
-  using DiffJ2OrbitalType = DiffTwoBodyJastrowOrbital<RadFuncType>;
-};
-
 #if defined(ENABLE_OFFLOAD)
 template<>
 class JastrowTypeHelper<BsplineFunctor<RadialJastrowBuilder::RealType>, RadialJastrowBuilder::detail::OMPTARGET>
 {
 public:
   using RadFuncType = BsplineFunctor<RadialJastrowBuilder::RealType>;
-  using J2OrbitalType = J2OMPTarget<RadFuncType>;
-  using DiffJ2OrbitalType = DiffTwoBodyJastrowOrbital<RadFuncType>;
+  using J2Type      = J2OMPTarget<RadFuncType>;
+  using DiffJ2Type  = DiffTwoBodyJastrowOrbital<RadFuncType>;
 };
 #endif
 
@@ -168,15 +160,15 @@ template<class RadFuncType, unsigned Implementation>
 std::unique_ptr<WaveFunctionComponent> RadialJastrowBuilder::createJ2(xmlNodePtr cur)
 {
   ReportEngine PRE(ClassName, "createJ2(xmlNodePtr)");
-  using Real                = typename RadFuncType::real_type;
-  using J2OrbitalType     = typename JastrowTypeHelper<RadFuncType, Implementation>::J2OrbitalType;
-  using DiffJ2OrbitalType = typename JastrowTypeHelper<RadFuncType, Implementation>::DiffJ2OrbitalType;
+  using Real       = typename RadFuncType::real_type;
+  using J2Type     = typename JastrowTypeHelper<RadFuncType, Implementation>::J2Type;
+  using DiffJ2Type = typename JastrowTypeHelper<RadFuncType, Implementation>::DiffJ2Type;
 
   XMLAttrString input_name(cur, "name");
   std::string j2name = input_name.empty() ? "J2_" + Jastfunction : input_name;
   SpeciesSet& species(targetPtcl.getSpeciesSet());
-  auto J2   = std::make_unique<J2OrbitalType>(j2name, targetPtcl);
-  auto dJ2  = std::make_unique<DiffJ2OrbitalType>(targetPtcl);
+  auto J2  = std::make_unique<J2Type>(j2name, targetPtcl);
+  auto dJ2 = std::make_unique<DiffJ2Type>(targetPtcl);
 
   std::string init_mode("0");
   {
@@ -348,17 +340,18 @@ std::unique_ptr<WaveFunctionComponent> RadialJastrowBuilder::createJ2<RPAFunctor
   return rpajastrow;
 }
 
-template<class RadFuncType, unsigned Implementation>
+template<class RadFuncType, bool SPIN, unsigned Implementation>
 std::unique_ptr<WaveFunctionComponent> RadialJastrowBuilder::createJ1(xmlNodePtr cur)
 {
   ReportEngine PRE(ClassName, "createJ1(xmlNodePtr)");
-  using Real            = typename RadFuncType::real_type;
-  using J1OrbitalType = typename JastrowTypeHelper<RadFuncType, Implementation>::J1OrbitalType;
+  using Real   = typename RadFuncType::real_type;
+  using TH     = JastrowTypeHelper<RadFuncType, Implementation>;
+  using J1Type = typename std::conditional<SPIN, typename TH::J1SpinType, typename TH::J1Type>::type;
 
   XMLAttrString input_name(cur, "name");
   std::string jname = input_name.empty() ? Jastfunction : input_name;
 
-  auto J1                = std::make_unique<J1OrbitalType>(jname, *SourcePtcl, targetPtcl);
+  auto J1 = std::make_unique<J1Type>(jname, *SourcePtcl, targetPtcl);
 
   xmlNodePtr kids = cur->xmlChildrenNode;
 
@@ -400,7 +393,8 @@ std::unique_ptr<WaveFunctionComponent> RadialJastrowBuilder::createJ1(xmlNodePtr
                       targetPtcl.getName(),
                   true);
       }
-      app_summary() << "    Radial function for element: " << speciesA << std::endl;
+      app_summary() << "    Radial function for element: " << speciesA << " - "
+                    << (speciesB.empty() ? targetPtcl.getName() : speciesB) << std::endl;
       functor->put(kids);
       app_summary() << std::endl;
       if (is_manager())
@@ -444,12 +438,12 @@ std::unique_ptr<WaveFunctionComponent> RadialJastrowBuilder::createJ1(xmlNodePtr
 template<>
 std::unique_ptr<WaveFunctionComponent> RadialJastrowBuilder::createJ1<RPAFunctor>(xmlNodePtr cur)
 {
-  using Real               = RealType;
+  using Real             = RealType;
   using SplineEngineType = CubicBspline<Real, LINEAR_1DGRID, FIRSTDERIV_CONSTRAINTS>;
   using RadFunctorType   = CubicSplineSingle<Real, SplineEngineType>;
   using GridType         = LinearGrid<Real>;
   using HandlerType      = LRHandlerBase;
-  using J1OrbitalType    = J1OrbitalSoA<RadFunctorType>;
+  using J1Type           = J1OrbitalSoA<RadFunctorType>;
 
   std::string input_name;
   std::string rpafunc = "RPA";
@@ -489,18 +483,18 @@ std::unique_ptr<WaveFunctionComponent> RadialJastrowBuilder::createJ1<RPAFunctor
   }
   myHandler->Breakup(targetPtcl, Rs);
 
-  Real Rcut          = myHandler->get_rc() - 0.1;
+  Real Rcut        = myHandler->get_rc() - 0.1;
   GridType* myGrid = new GridType;
   int npts         = static_cast<int>(Rcut / 0.01) + 1;
   myGrid->set(0, Rcut, npts);
 
   //create the numerical functor
-  auto nfunc                     = std::make_unique<RadFunctorType>();
+  auto nfunc                       = std::make_unique<RadFunctorType>();
   ShortRangePartAdapter<Real>* SRA = new ShortRangePartAdapter<Real>(myHandler);
   SRA->setRmax(Rcut);
   nfunc->initialize(SRA, myGrid);
 
-  auto J1                = std::make_unique<J1OrbitalType>(jname, *SourcePtcl, targetPtcl);
+  auto J1 = std::make_unique<J1Type>(jname, *SourcePtcl, targetPtcl);
 
   SpeciesSet& sSet = SourcePtcl->getSpeciesSet();
   for (int ig = 0; ig < sSet.getTotalNum(); ig++)
@@ -521,7 +515,7 @@ std::unique_ptr<WaveFunctionComponent> RadialJastrowBuilder::buildComponent(xmlN
   aAttrib.add(NameOpt, "name");
   aAttrib.add(TypeOpt, "type");
   aAttrib.add(Jastfunction, "function");
-  aAttrib.add(SpinOpt, "spin");
+  aAttrib.add(SpinOpt, "spin", {"no", "yes"});
 #if defined(ENABLE_OFFLOAD)
   aAttrib.add(useGPU, "gpu", {"yes", "no"});
 #endif
@@ -539,16 +533,30 @@ std::unique_ptr<WaveFunctionComponent> RadialJastrowBuilder::buildComponent(xmlN
     // it's a one body jastrow factor
     if (Jastfunction == "bspline")
     {
+      if (SpinOpt == "yes")
+      {
 #if defined(QMC_CUDA)
-      return createJ1<BsplineFunctor<RealType>, detail::CUDA_LEGACY>(cur);
+        myComm->barrier_and_abort("RadialJastrowBuilder::buildComponent spin resolved bspline Jastrow is not supported in legacy CUDA build.");
 #else
-      return createJ1<BsplineFunctor<RealType>>(cur);
+        return createJ1<BsplineFunctor<RealType>, true>(cur);
 #endif
+      }
+      else
+      {
+#if defined(QMC_CUDA)
+        return createJ1<BsplineFunctor<RealType>, false, detail::CUDA_LEGACY>(cur);
+#else
+        return createJ1<BsplineFunctor<RealType>>(cur);
+#endif
+      }
     }
     else if (Jastfunction == "pade")
     {
       guardAgainstPBC();
-      return createJ1<PadeFunctor<RealType>>(cur);
+      if (SpinOpt == "yes")
+        return createJ1<PadeFunctor<RealType>, true>(cur);
+      else
+        return createJ1<PadeFunctor<RealType>>(cur);
     }
     else if (Jastfunction == "pade2")
     {
@@ -558,11 +566,17 @@ std::unique_ptr<WaveFunctionComponent> RadialJastrowBuilder::buildComponent(xmlN
     else if (Jastfunction == "shortrangecusp")
     {
       //guardAgainstPBC(); // is this needed?
-      return createJ1<ShortRangeCuspFunctor<RealType>>(cur);
+      if (SpinOpt == "yes")
+        return createJ1<ShortRangeCuspFunctor<RealType>, true>(cur);
+      else
+        return createJ1<ShortRangeCuspFunctor<RealType>>(cur);
     }
     else if (Jastfunction == "user")
     {
-      return createJ1<UserFunctor<RealType>>(cur);
+      if (SpinOpt == "yes")
+        return createJ1<UserFunctor<RealType>, true>(cur);
+      else
+        return createJ1<UserFunctor<RealType>>(cur);
     }
     else if (Jastfunction == "rpa")
     {
@@ -587,13 +601,14 @@ std::unique_ptr<WaveFunctionComponent> RadialJastrowBuilder::buildComponent(xmlN
 #if defined(ENABLE_OFFLOAD)
       if (useGPU == "yes")
       {
-        static_assert(std::is_same<JastrowTypeHelper<BsplineFunctor<RealType>, OMPTARGET>::J2OrbitalType,
-                                   J2OMPTarget<BsplineFunctor<RealType>>>::value, "check consistent type");
-        if(targetPtcl.getCoordinates().getKind() != DynamicCoordinateKind::DC_POS_OFFLOAD)
+        static_assert(std::is_same<JastrowTypeHelper<BsplineFunctor<RealType>, OMPTARGET>::J2Type,
+                                   J2OMPTarget<BsplineFunctor<RealType>>>::value,
+                      "check consistent type");
+        if (targetPtcl.getCoordinates().getKind() != DynamicCoordinateKind::DC_POS_OFFLOAD)
         {
           std::ostringstream msg;
-          msg << "Offload enabled Jastrow needs the gpu=\"yes\" attribute in the \""
-              << targetPtcl.getName() << "\" particleset" << std::endl;
+          msg << "Offload enabled Jastrow needs the gpu=\"yes\" attribute in the \"" << targetPtcl.getName()
+              << "\" particleset" << std::endl;
           myComm->barrier_and_abort(msg.str());
         }
         app_summary() << "    Running on an accelerator via OpenMP offload." << std::endl;

src/QMCWaveFunctions/Jastrow/RadialJastrowBuilder.h

@@ -61,7 +61,7 @@ class RadialJastrowBuilder : public WaveFunctionComponentBuilder
   ParticleSet* SourcePtcl;
 
   // has a specialization for RPAFunctor in cpp file
-  template<class RadFuncType, unsigned Implementation = detail::CPU>
+  template<class RadFuncType, bool SPIN = false, unsigned Implementation = detail::CPU>
   std::unique_ptr<WaveFunctionComponent> createJ1(xmlNodePtr cur);
 
   template<class RadFuncType, unsigned Implementation = detail::CPU>

src/QMCWaveFunctions/tests/CMakeLists.txt

@@ -93,6 +93,7 @@ set(JASTROW_SRC
     test_pade_jastrow.cpp
     test_short_range_cusp_jastrow.cpp
     test_J1OrbitalSoA.cpp
+    test_J1Spin.cpp
     test_J2_bspline.cpp
     test_DiffTwoBodyJastrowOrbital.cpp)
 set(DETERMINANT_SRC

src/QMCWaveFunctions/tests/test_J1Spin.cpp

@@ -0,0 +1,115 @@
+//////////////////////////////////////////////////////////////////////////////////////
+// This file is distributed under the University of Illinois/NCSA Open Source License.
+// See LICENSE file in top directory for details.
+//
+// Copyright (c) 2021 QMCPACK developers.
+//
+// File developed by: Shiv Upadhyay, shivnupadhyay@gmail.com, University of Pittsburgh
+//
+// File created by: Shiv Upadhyay, shivnupadhyay@gmail.com, University of Pittsburgh
+//////////////////////////////////////////////////////////////////////////////////////
+
+#include "catch.hpp"
+
+#include "Particle/ParticleSetPool.h"
+#include "QMCWaveFunctions/Jastrow/J1Spin.h"
+#include "QMCWaveFunctions/Jastrow/RadialJastrowBuilder.h"
+#include "QMCWaveFunctions/WaveFunctionFactory.h"
+
+namespace qmcplusplus
+{
+#if !defined(QMC_CUDA)
+TEST_CASE("J1 spin evaluate derivatives Jastrow", "[wavefunction]")
+{
+  Communicate* c = OHMMS::Controller;
+  auto ions_uptr = std::make_unique<ParticleSet>();
+  auto elec_uptr = std::make_unique<ParticleSet>();
+  ParticleSet& ions_(*ions_uptr);
+  ParticleSet& elec_(*elec_uptr);
+
+  ions_.setName("ion0");
+  ions_.create(1);
+  ions_.R[0]                 = {0.0, 0.0, 0.0};
+  SpeciesSet& ispecies       = ions_.getSpeciesSet();
+  int HIdx                   = ispecies.addSpecies("H");
+  int ichargeIdx             = ispecies.addAttribute("charge");
+  ispecies(ichargeIdx, HIdx) = 1.0;
+
+  elec_.setName("e");
+  elec_.create({1, 1});
+  elec_.R[0] = {0.5, 0.5, 0.5};
+  elec_.R[1] = {-0.5, -0.5, -0.5};
+
+  SpeciesSet& tspecies       = elec_.getSpeciesSet();
+  int upIdx                  = tspecies.addSpecies("u");
+  int downIdx                = tspecies.addSpecies("d");
+  int massIdx                = tspecies.addAttribute("mass");
+  int chargeIdx              = tspecies.addAttribute("charge");
+  tspecies(massIdx, upIdx)   = 1.0;
+  tspecies(massIdx, downIdx) = 1.0;
+  tspecies(chargeIdx, upIdx) = -1.0;
+  tspecies(massIdx, downIdx) = -1.0;
+  // Necessary to set mass
+  elec_.resetGroups();
+
+  ParticleSetPool ptcl = ParticleSetPool(c);
+  ptcl.addParticleSet(std::move(elec_uptr));
+  ptcl.addParticleSet(std::move(ions_uptr));
+
+
+  ions_.update();
+  elec_.addTable(elec_);
+  elec_.addTable(ions_);
+  elec_.update();
+
+  const char* jasxml = "<wavefunction name=\"psi0\" target=\"e\"> \
+  <jastrow name=\"J1\" type=\"One-Body\" function=\"Bspline\" print=\"yes\" source=\"ion0\" spin=\"yes\"> \
+    <correlation speciesA=\"H\" speciesB=\"u\" cusp=\"0.0\" size=\"2\" rcut=\"5.0\"> \
+      <coefficients id=\"J1uH\" type=\"Array\"> 0.5 0.1 </coefficients> \
+    </correlation> \
+    <correlation speciesA=\"H\" speciesB=\"d\" cusp=\"0.0\" size=\"2\" rcut=\"5.0\"> \
+      <coefficients id=\"J1dH\" type=\"Array\"> 0.5 0.1 </coefficients> \
+    </correlation> \
+  </jastrow> \
+</wavefunction> \
+";
+  Libxml2Document doc;
+  bool okay = doc.parseFromString(jasxml);
+  REQUIRE(okay);
+
+  xmlNodePtr jas1 = doc.getRoot();
+
+  // update all distance tables
+  elec_.update();
+  WaveFunctionFactory wf_factory("psi0", elec_, ptcl.getPool(), c);
+  wf_factory.put(jas1);
+  auto& twf(*wf_factory.getTWF());
+  twf.setMassTerm(elec_);
+  twf.evaluateLog(elec_);
+  twf.prepareGroup(elec_, 0);
+
+  auto& twf_component_list = twf.getOrbitals();
+
+  opt_variables_type active;
+  twf.checkInVariables(active);
+  active.removeInactive();
+  int nparam = active.size_of_active();
+  REQUIRE(nparam == 4);
+
+  using ValueType = QMCTraits::ValueType;
+  std::vector<ValueType> dlogpsi(nparam);
+  std::vector<ValueType> dhpsioverpsi(nparam);
+  //twf.evaluateDerivatives(elec_, active, dlogpsi, dhpsioverpsi);
+  twf_component_list[0]->evaluateDerivatives(elec_, active, dlogpsi, dhpsioverpsi);
+
+  // Numbers not validated
+  std::vector<ValueType> expected_dlogpsi      = {-0.46681472435, -0.5098025897, -0.46681472435, -0.5098025897};
+  std::vector<ValueType> expected_dhpsioverpsi = {-0.5798216548, 0.37977462695, -0.5798216548, 0.37977462695};
+  for (int i = 0; i < nparam; i++)
+  {
+    CHECK(dlogpsi[i] == ValueApprox(expected_dlogpsi[i]));
+    CHECK(dhpsioverpsi[i] == ValueApprox(expected_dhpsioverpsi[i]));
+  }
+}
+#endif
+} // namespace qmcplusplus

tests/molecules/H2_ae/CMakeLists.txt

@@ -19,6 +19,18 @@ if(NOT QMC_CUDA)
         H2_OPT_SCALARS # Final VMC step
         TRUE)
 
+      qmc_run_and_check(
+        short-H2-orb-opt-spindepjas
+        "${qmcpack_SOURCE_DIR}/tests/molecules/H2_ae"
+        H2
+        h2_orb_opt_spindepjas.xml
+        16
+        1
+        ${SUCCESS_STATUS_MP}
+        10
+        H2_OPT_SCALARS # Final VMC step
+        TRUE)
+
     else()
       message(VERBOSE "Skipping H2 test because lmyengine interface was not built (BUILD_LMYENGINE_INTERFACE=0)")
     endif(BUILD_LMYENGINE_INTERFACE)