Skip to content

Commit

Permalink
Merge pull request #3257 from shivupa/spinonebody
Browse files Browse the repository at this point in the history
One Body Spin dependent Jastrow
  • Loading branch information
ye-luo authored Aug 30, 2021
2 parents 9818a47 + 897da66 commit 6161091
Show file tree
Hide file tree
Showing 8 changed files with 1,471 additions and 48 deletions.
744 changes: 744 additions & 0 deletions src/QMCWaveFunctions/Jastrow/J1Spin.h

Large diffs are not rendered by default.

109 changes: 62 additions & 47 deletions src/QMCWaveFunctions/Jastrow/RadialJastrowBuilder.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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

Expand Down Expand Up @@ -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");
{
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -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";
Expand Down Expand Up @@ -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++)
Expand All @@ -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
Expand All @@ -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")
{
Expand All @@ -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")
{
Expand All @@ -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;
Expand Down
2 changes: 1 addition & 1 deletion src/QMCWaveFunctions/Jastrow/RadialJastrowBuilder.h
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down
1 change: 1 addition & 0 deletions src/QMCWaveFunctions/tests/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
115 changes: 115 additions & 0 deletions src/QMCWaveFunctions/tests/test_J1Spin.cpp
Original file line number Diff line number Diff line change
@@ -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
12 changes: 12 additions & 0 deletions tests/molecules/H2_ae/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading

0 comments on commit 6161091

Please sign in to comment.